/*
 *  cudadl version 0.9: Compute discrete logs in smooth group orders
 *                      using CUDA
 *  Copyright (C) 2012 by Ryan Henry and Ian Goldberg
 *                        {rhenry,iang}@cs.uwaterloo.ca
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of version 3 of the GNU General Public License as
 *  published by the Free Software Foundation.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <fstream>
#include <unistd.h>
#include <sys/time.h>
#include <NTL/vec_ZZ.h>
#include <NTL/ZZ_p.h>

#include <map>
#include <sstream>

#include <sys/types.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <cuda_runtime.h>

#include "atomic_iostream.h"
#include "desired_resources.h"
#include "cudadl.h"

NTL_CLIENT

string output_prefix;

#ifdef SAVE_DPS
static ofstream dp_file_stream;
#endif

typedef map<std::string, pair<ZZ,ZZ> > DTable;

struct CBData {
    const ZZ_p &base;
    const ZZ_p &target;
    const ZZ &order;
    unsigned long long numdp;
    DTable dtable;
    bool found_collision;
    ZZ expon;

    CBData(const ZZ_p &_base, const ZZ_p &_target, const ZZ &_order) :
	base(_base), target(_target), order(_order), numdp(0),
	found_collision(false) {}
};

// This function is called from inside cuda_dl for each DP it encounters.
// It calls the function named "dpcallback" directly.  It would be
// cleaner if this were passed as a function pointer to cuda_dl, but
// that makes nvcc 3.1 segfault. :-p
// dp points to an array of WORDS+7 unsigned ints:
// - 1 word of threadID/blockID
// - WORDS words of the dp value
// - 3 words of a
// - 3 words of b
bool dpcallback(void *cbdata, unsigned int *dpwords)
{
    CBData *d = (CBData*)cbdata;
    // WARNING: this assumes
    //    sizeof(unsigned long) == sizeof(unsigned long long) !
    ZZ zz_a = to_ZZ(dpwords[WORDS+3]);
    zz_a <<= 32;
    zz_a += dpwords[WORDS+2];
    zz_a <<= 32;
    zz_a += dpwords[WORDS+1];
    ZZ zz_b = to_ZZ(dpwords[WORDS+6]);
    zz_b <<= 32;
    zz_b += dpwords[WORDS+5];
    zz_b <<= 32;
    zz_b += dpwords[WORDS+4];

    //ZZ_p dp = power(d->base, zz_a) * power(d->target, zz_b);

    // cerr << "DP " << ++(d->numdp) << "\r";
    pair<ZZ,ZZ> ab(zz_a,zz_b);
    pair<DTable::iterator, bool> res;
    string x((const char *)(dpwords+1), WORDS*sizeof(unsigned int));

#ifdef SAVE_DPS
    if (!d->found_collision) {
	ZZ zz_x;
	ZZFromBytes(zz_x, (const unsigned char *)(dpwords+1),
		    WORDS*sizeof(unsigned int));
	dp_file_stream << zz_x << "\n";
    }
#endif

    res = d->dtable.insert(DTable::value_type(x, ab));
    if (!res.second) {
	// Collision!
	ZZ adiff = to_ZZ(res.first->second.first) - zz_a;
	ZZ bdiff = zz_b - to_ZZ(res.first->second.second);
	if (bdiff < 0) bdiff += d->order;
	ZZ binv;
	if (InvModStatus(binv, bdiff, d->order) == 0) {
	    d->expon = MulMod(binv, adiff, d->order);
	    d->found_collision = true;
	} else {
	    if (d->order > (1<<20)) {
		cerr << "Unhelpful collision\n";
	    }
	}
    }

    return d->found_collision;
}

// Compute the discrete log of target mod p, to the given base.
// p must be the current ZZ_p modulus.
// Place the result in exp.  fvec is a vector of the factors of
// (p-1)/2, which must each be small enough to compute discrete logs
// with some other method (kangaroo, index calculus, GNFS, etc.).
// label is "p" or "q", to be printed to report progress.
// Return 0 on failure, 1 on success.
static int p_dl(const ZZ_p &target, const ZZ_p &base, ZZ &exp,
    const vec_ZZ &fvec, const string &label, unsigned int initial_subproblem_id,
    unsigned short GB_mem)
{
    const int flen = fvec.length();

    // Compute phi(p)
    ZZ phip;
    phip = 2;
    for (int i = 0; i < flen; ++i) {
	phip *= fvec[i];
    }

    // Invariant: the desired exp \equiv (curexp mod curmodulus)
    ZZ curmodulus, curexp;
    curmodulus = 2;
    curexp = 0;

    for (int i = 0; i < flen; ++i) {
	unsigned short desired_dpnodes;
	unsigned int max_workers;
	unsigned int dpfreq;

	unsigned short freq_reduction_threshold = 4;
	// this should be okay in most cases and shouldn't cause
	// any negative effects if it's too low

	desired_resources(fvec[i], 1, 1, GB_mem, freq_reduction_threshold, desired_dpnodes, max_workers, dpfreq);

	ZZ quotient;
	ZZ_p subgroup_base;
	ZZ_p subgroup_target;
	
	{
	    AtomicWriter atomic_cout(cout);
	    atomic_cout << label << " submodulus " << i+1 << " of " << flen << "...\n";
	    atomic_cout << "Settings (memory, dpfreq): " << GB_mem << ", " << dpfreq << "\n";
	    // Figure out exp mod fvec[i] by taking each side to the power
	    // of phirho/fvec[i] so that we're working in the
	    // subgroup of order fvec[i].
	    quotient = phip / fvec[i];
	    subgroup_base = power(base, quotient);
	    subgroup_target = power(target, quotient);

	    if (subgroup_base == 1) {
		// The original base wasn't a generator of the whole group.
		if (subgroup_target == 1) {
		    atomic_cout << "Non-unique solution (mod " << fvec[i] <<")\n";
		    continue;
		} else {
		    atomic_cout << "Target not in subgroup generated by base\n";
		    return 0;
		}
	    }
	}

#ifdef SAVE_DPS
	std::ostringstream oss;
	oss << "dplist_" << i+initial_subproblem_id << ".out";
	dp_file_stream.open(oss.str().c_str());
#endif

	// Now use your favourite method to get the DL of
	// subgroup_target with base subgroup_base, knowing that it's in
	// the range [0,fvec[i]).

	ZZ md = ZZ_p::modulus();
	CBData cbdata(subgroup_base, subgroup_target, fvec[i]);
	struct timeval st, et;

#ifdef SAVE_DPS
	dp_file_stream << "Subproblem " << i+initial_subproblem_id << "\n";
#endif
	gettimeofday(&st, NULL);
	unsigned int launch_count = 0;
	bool filled_dp_buffer = false;
	{
#ifdef DERANDOMIZE
	    RandomStreamPush push_seed;
	    // the seed will be reset to its original value
	    // once we exit this scope
	    SetSeed(rep(subgroup_base)*rep(subgroup_target)*fvec[i]*md);
#endif
	    cuda_dl(subgroup_base, subgroup_target, fvec[i], md, dpfreq, &cbdata, &launch_count, &filled_dp_buffer);
	}
	ZZ subgroup_dl = cbdata.expon;
	gettimeofday(&et, NULL);
	unsigned long us_elapsed = (et.tv_sec-st.tv_sec)*1000000 +
	    (et.tv_usec-st.tv_usec);
	printf("%ld.%06ld seconds elapsed\n", us_elapsed/1000000,
	    us_elapsed % 1000000);

	AtomicWriter(cout) << "Timing (subproblemid, label, launches): "
	                   << i+initial_subproblem_id << ", "
	                   << label << ", "
	                   << launch_count << "\n" << std::flush;

	CRT(curexp, curmodulus, subgroup_dl, fvec[i]);
	// cout << "CRT\n";
	// cout << "curexp = " << curexp << "\n";
	// cout << "curmodulus = " << curmodulus << "\n\n";

#ifdef SAVE_DPS
	if (filled_dp_buffer) {
	    AtomicWriter atomic_cout(cout);
	    atomic_cout << "Warning: The device dp buffer was filled, so some points were not recorded.\n";
	    atomic_cout << "These points will not be reproducible (subproblem " << i+initial_subproblem_id << ").\n";
	}
        dp_file_stream.close();
#endif
    }
    // We'd like a non-negative answer back
    if (curexp >= 0) {
	exp = curexp;
    } else {
	exp = curexp + curmodulus;
    }
    return 1;
}

typedef struct {
    pid_t pid;
    int rfd;
} PDLHandle;

// Behave like p_dl, but do the work in an asynchronous subprocess
static PDLHandle* p_dl_fork_start(const ZZ &p, const ZZ_p &target,
    const ZZ_p &base, const vec_ZZ &fvec, const string &label,
    int deviceid, unsigned int initial_subproblem_id,
    unsigned short GB_mem_per_subprocess)
{
    PDLHandle *handle = new PDLHandle;
    int fds[2];
    int res = socketpair(AF_UNIX, SOCK_STREAM, 0, fds);
    if (res < 0) {
	perror("socketpair");
	delete handle;
	return NULL;
    }
    handle->rfd = fds[0];

    cout.flush();
    cerr.flush();
    // flush the output before forking

    pid_t childpid = fork();
    if (childpid == -1) {
	perror("fork");
	delete handle;
	return NULL;
    }
    if (childpid == 0) {
	// Child
	close(fds[0]);
	int wfd = fds[1];

	ZZ exp;
	ZZ_p::init(p);
	ZZ_p target_p, base_p;
	conv(target_p, rep(target));
	conv(base_p, rep(base));

	cudaError_t cudares = cudaSetDevice(deviceid);
	if (cudares != cudaSuccess) {
	    AtomicWriter(cerr) << "Error setting CUDA device: " << cudaGetErrorString(cudares) << "\n";
	    exit(1);
	}
	int res = p_dl(target_p, base_p, exp, fvec, label, initial_subproblem_id, GB_mem_per_subprocess);

	if (res) {
	    // Write the result back to the parent using wfd
	    unsigned short explen = NumBytes(exp);
	    res = write(wfd, &explen, sizeof(unsigned short));
	    unsigned char expbuf[explen];
	    BytesFromZZ(expbuf, exp, explen);
	    res = write(wfd, expbuf, explen);
	    close(wfd);
	}

	exit(0);
    } else {
	// Parent
	close(fds[1]);
    }

    return handle;
}

// Wait until the subprocess started by p_dl_fork_start completes, and
// return its result.  handle is cleaned up.
static int p_dl_fork_join(PDLHandle *handle, ZZ &exp)
{
    if (!handle) return 0;
    pid_t childpid = handle->pid;
    unsigned short explen;
    int res = read(handle->rfd, &explen, sizeof(unsigned short));
    if (res < (int)sizeof(unsigned short)) {
	close(handle->rfd);
	delete handle;
	waitpid(childpid, NULL, 0);
	return 0;
    }
    unsigned char expbuf[explen];
    res = read(handle->rfd, expbuf, explen);
    if (res < explen) {
	close(handle->rfd);
	delete handle;
	waitpid(childpid, NULL, 0);
	return 0;
    }
    ZZFromBytes(exp, expbuf, explen);
    close(handle->rfd);
    delete handle;
    waitpid(childpid, NULL, 0);
    return 1;
}

int main(int argc, char **argv)
{
#ifdef SAVE_DPS
    cerr << "Note: Saving the distinguished points. Do not run huge problems or else it will use up all of your disk space.\n";
#ifndef DERANDOMIZE
    cerr << "Saving DPs without derandomization!\n";
#endif
#endif

    if (argc != 2) {
	cerr << "Usage: " << argv[0] << " total_mem_GB";
	exit(1);
    }

    unsigned short GB_mem_per_node = strtol(argv[1], NULL, 10);

    // Initialize the prng with some randomness from the kernel
    unsigned char randbuf[1024];
    ifstream urand("/dev/urandom");
    urand.read((char *)randbuf, sizeof(randbuf));
    urand.close();
    ZZ randzz = ZZFromBytes(randbuf, sizeof(randbuf));
    SetSeed(randzz);

    ZZ rho, p, q;
    vec_ZZ pfvec, qfvec;
    cin >> rho >> p >> pfvec >> q >> qfvec;

    // Generate a DLP mod rho (in the large odd-order subgroup)
    ZZ_p::init(rho);
    ZZ_p base;
    ZZ_p target;
    {
#ifdef DERANDOMIZE
	RandomStreamPush push_seed;
	// the seed will be reset to its original value
	// once we exit this scope
	SetSeed(rho*p*q);
#endif
	base = power(random_ZZ_p(), 2);
	target = power(random_ZZ_p(), 2);
    }

    cout << "base = " << base << "\n";
    cout << "target = " << target << "\n";

    ZZ exp_p, exp_q, exp;
    int res_p = 0;
    int res_q = 0;

    PDLHandle *handle_p, *handle_q;

    handle_p = p_dl_fork_start(p, target, base, pfvec, "p", 0, 0, GB_mem_per_node/2);
    handle_q = p_dl_fork_start(q, target, base, qfvec, "q", 1, pfvec.length(), GB_mem_per_node/2);
    res_p = p_dl_fork_join(handle_p, exp_p);
    res_q = p_dl_fork_join(handle_q, exp_q);

    if (res_p && res_q) {
	ZZ_p::init(rho);
	ZZ pm1 = (p - 1)/2;
	ZZ qm1 = (q - 1)/2;
	if (exp_p < 0) exp_p += pm1;
	if (exp_q < 0) exp_q += qm1;
	CRT(exp_p, pm1, exp_q, qm1);
	exp = exp_p;
	// We'd like a non-negative answer back
	if (exp < 0) exp += pm1;

	cout << "exp = " << exp << "\n";
	ZZ_p base_exp;
	power(base_exp, base, exp);
	if (base_exp == target) {
	    cout << "CORRECT!\n";
	} else {
	    cout << "INCORRECT:\nbase^exp = " << base_exp << "\n";
	    cout << "target   = " << target << "\n";
	}
    } else {
	cout << "FAIL\n";
    }

    return 0;
}