// $Id: TestGenerator.cc 379 2007-09-04 18:24:09Z dflater $

#include <iostream>
#include <sstream>
#include <fstream>
#include <string>
#include <vector>
#include <map>
#include <algorithm>
#include <ext/algorithm>
#include <stdlib.h>
#include <assert.h>
#include "Contest.hh"
#include "misc.hh"
#include "parser.hh"


void usage () {
  std::cerr << "Usage:  TestGenerator input-filename > output.sql\n"
    "See Documentation.pdf for format of input.\n";
  exit (-1);
}


void makeContest (ContestVector &contests,
                  const TGParser::Contest &contest,
		  unsigned long &contestSerial,
		  unsigned long &choiceSerial) {
  if (contest.level == TGParser::Contest::systemExtent)
    contests.push_back (Contest (contest, 0, contestSerial, choiceSerial));
  else {
    const unsigned long numDistrictsOrPrecincts (
      contest.level == TGParser::Contest::district ? election.districts
                                                   : election.precincts);
    for (unsigned long i=1; i<=numDistrictsOrPrecincts; ++i)
      contests.push_back (Contest (contest, i, contestSerial, choiceSerial));
  }
}


void makeReportingContexts (std::vector<std::string> &reportingContexts) {
  for (unsigned long i=1; i<=election.precincts; ++i)
    reportingContexts.push_back ("Precinct " + stringify(i));
  for (unsigned long i=1; i<=election.districts; ++i)
    reportingContexts.push_back ("District " + stringify(i));
  reportingContexts.push_back ("Jurisdiction");
}


void generateReportingContexts (const std::vector<std::string> &reportingContexts) {
  std::cout << "copy ReportingContext from stdin CSV;\n";
  for (std::vector<std::string>::const_iterator it (reportingContexts.begin());
       it != reportingContexts.end();
       ++it)
    std::cout << *it << '\n';
  std::cout << "\\.\n\n";
}


void generateContests (const ContestVector &contests) {
  std::cout << "copy Contest from stdin CSV;\n";
  for (ContestVector::const_iterator it (contests.begin());
       it != contests.end();
       ++it)
    std::cout << it->id
	      << ','
	      << it->description
              << ','
              << stringify(it->logic)
              << ','
              << it->N
              << ','
              << (it->W > 0 ? it->N : 0)
              << ",false\n";
  std::cout << "\\.\n\n";
}


void generateChoices (const ContestVector &contests) {
  std::cout << "copy Choice from stdin CSV;\n";
  for (ContestVector::const_iterator it (contests.begin());
       it != contests.end();
       ++it) {
    unsigned long choiceCtr (1);
    for (unsigned long i (it->firstChoice);
	 i<=it->lastChoice;
	 ++choiceCtr, ++i)
      std::cout << i
		<< ','
		<< it->id
		<< ",Choice #"
		<< choiceCtr
		<< " in contest #"
		<< it->id
		<< ",,"
                << (choiceCtr <= it->W ? "true" : "false")
                << '\n';
  }
  std::cout << "\\.\n\n";
}


void makePrecincts (std::vector<unsigned long> &districtMap_out) {
  districtMap_out.resize (election.precincts+1);
  for (unsigned long i=1; i<=election.precincts; ++i) {
    switch (election.precinctDistribution) {
    case TGParser::Election::uniformRandom:
      districtMap_out[i] = randFunc(election.districts)+1;
      break;
    case TGParser::Election::even:
      districtMap_out[i] = evenDist(i,election.precincts,election.districts);
      break;
    default:
      assert (false);
    }
  }
}


void makeBallots (std::vector<unsigned long> &styleMap_out) {
  styleMap_out.resize (election.ballots+1);
  for (unsigned long i=1; i<=election.ballots; ++i) {
    switch (election.ballotDistribution) {
    case TGParser::Election::uniformRandom:
      styleMap_out[i] = randFunc(election.precincts)+1;
      break;
    case TGParser::Election::even:
      styleMap_out[i] = evenDist(i,election.ballots,election.precincts);
      break;
    default:
      assert (false);
    }
  }
}


void generateBallotStyles (const ContestVector &contests,
			   const std::vector<unsigned long> &districtMap) {
  std::cout << "copy BallotStyle from stdin CSV;\n";
  for (unsigned long i=1; i<=election.precincts; ++i)
    std::cout << i
	      << ",Precinct "
	      << i
	      << " style\n";
  std::cout << "\\.\n\n";

  std::cout << "copy BallotStyleContestAssociation from stdin CSV;\n";
  for (unsigned long i=1; i<=election.precincts; ++i)
    for (ContestVector::const_iterator it (contests.begin());
	 it != contests.end();
         ++it)
      if (pertinent (*it, districtMap[i], i))
        std::cout << i
		  << ','
		  << it->id
		  << '\n';
  std::cout << "\\.\n\n";

  std::cout << "copy BallotStyleReportingContextAssociation from stdin CSV;\n";
  for (unsigned long i=1; i<=election.precincts; ++i)
    std::cout << i
	      << ",Precinct "
	      << i
	      << '\n'
	      << i
	      << ",District "
	      << districtMap[i]
	      << '\n'
	      << i
	      << ",Jurisdiction\n";
  std::cout << "\\.\n\n";
}


void generateBallots (const std::vector<unsigned long> &styleMap) {
  std::cout << "copy Ballot from stdin CSV;\n";
  for (unsigned long i=1; i<=election.ballots; ++i)
    std::cout << i
	      << ','
	      << styleMap[i]
	      << ",true\n";
  std::cout << "\\.\n\n";
}


// A chooser function returns a vector representing the votes that one
// ballot casts in one contest, according to some pattern / policy.
//   N-of-M:  The choices voted for are simply listed.
//   Cumulative:  If a choice gets more than one vote, it is listed
//                more than once.
//   Ranked order:  Choices are listed in order of decreasing preference.

typedef std::vector<unsigned long> (*ChooserFn) (Contest &contest);


// FIXME, the protocol for returning cumulative votes no longer
// simplifies anything.
std::vector<unsigned long> triangleChooser (Contest &contest) {
  std::vector<unsigned long> choices;
  if (contest.logic == TGParser::Contest::cumulative) {
    assert (contest.lastChoice >= contest.firstChoice);
    const unsigned long M (contest.lastChoice - contest.firstChoice + 1);
    for (unsigned long i(0); i<M && choices.size()<contest.N; ) {
      if (contest.tallies[i] < (i+1)*contest.triangleMultiplier) {
	choices.push_back (i + contest.firstChoice);
	++(contest.tallies[i]);
      } else
        ++i;
    }

  } else if (contest.logic == TGParser::Contest::NofM) {
    // While complicated, this approach minimizes wasted votes.
    assert (contest.lastChoice >= contest.firstChoice);
    const unsigned long M (contest.lastChoice - contest.firstChoice + 1);
    while (choices.size() < contest.N) {
      unsigned long mostEligibleChoice (0), mostEligibleDelta (0);
      for (unsigned long i(0); i<M; ++i)
        if (contest.tallies[i] < (i+1)*contest.triangleMultiplier &&
	    find (choices.begin(), choices.end(), i+contest.firstChoice)
	    == choices.end()) {
	  const unsigned long delta ((i+1)*contest.triangleMultiplier
				     - contest.tallies[i]);
	  if (delta > mostEligibleDelta)
	    mostEligibleChoice = i, mostEligibleDelta = delta;
        }
      if (mostEligibleDelta > 0) {
	choices.push_back (mostEligibleChoice + contest.firstChoice);
	++(contest.tallies[mostEligibleChoice]);
      } else
        break;
    }

  } else {
    assert (contest.logic == TGParser::Contest::rankedOrder);
    assert (contest.firstChoice > 0);
    for (unsigned long i (contest.lastChoice); i>=contest.firstChoice; --i)
      choices.push_back (i);
  }
  return choices;
}


std::vector<unsigned long> uniformRandomChooser (Contest &contest) {
  std::vector<unsigned long> ballotPositions;
  for (unsigned long i (contest.firstChoice); i<=contest.lastChoice; ++i)
    ballotPositions.push_back (i);

  if (contest.logic == TGParser::Contest::rankedOrder) {
    std::random_shuffle (ballotPositions.begin(),
			 ballotPositions.end(), randFunc);
    return ballotPositions;

  } else if (contest.logic == TGParser::Contest::cumulative) {
    std::vector<unsigned long> votes;
    for (unsigned long i(0); i<contest.N; ++i)
      votes.push_back (ballotPositions[randFunc(ballotPositions.size())]);
    return votes;

  } else if (contest.logic == TGParser::Contest::NofM) {
    std::vector<unsigned long> votes (contest.N);
    __gnu_cxx::random_sample (ballotPositions.begin(),
			      ballotPositions.end(),
			      votes.begin(),
			      votes.end(),
			      randFunc);
    return votes;

  } else
    assert (false);
}


void doContest (unsigned long ballot, Contest &contest) {
  ChooserFn chooser;
  switch (contest.distribution) {
  case TGParser::Contest::uniformRandom:
    chooser = uniformRandomChooser;
    break;
  case TGParser::Contest::triangle:
    chooser = triangleChooser;
    break;
  default:
    assert (false);
  }
  std::vector<unsigned long> votes (chooser(contest));
  if (contest.logic == TGParser::Contest::rankedOrder) {
    for (unsigned long i (0); i<votes.size(); ++i)
      std::cout << ballot
		<< ','
		<< votes[i]
		<< ','
		<< i+1
		<< '\n';
  } else {
    // The following logic works for both N-of-M and cumulative,
    // though for N-of-M it does unnecessary work.
    std::map<unsigned long, unsigned long> combinedVotes;
    for (std::vector<unsigned long>::iterator it (votes.begin());
	 it != votes.end();
	 ++it) {
      std::map<unsigned long, unsigned long>::iterator it2
        (combinedVotes.find(*it));
      if (it2 == combinedVotes.end())
        combinedVotes[*it] = 1;
      else
	++(it2->second);
    }
    for (std::map<unsigned long, unsigned long>::iterator it
	   (combinedVotes.begin());
	 it != combinedVotes.end();
	 ++it)
      std::cout << ballot
		<< ','
		<< it->first
		<< ','
		<< it->second
		<< '\n';
  }
}


// After we know how many votes are available, figure out how much to
// scale up the Triangle distribution.

void figureMultipliers (ContestVector &contests,
			const std::vector<unsigned long> &districtMap,
			const std::vector<unsigned long> &styleMap) {
  for (ContestVector::iterator it (contests.begin());
       it != contests.end();
       ++it)
    if (it->distribution == TGParser::Contest::triangle) {
      unsigned long numBallots (0);
      for (unsigned long ballot (1); ballot <= election.ballots; ++ballot) {
	const unsigned long precinct (styleMap[ballot]);
	const unsigned long district (districtMap[precinct]);
	if (pertinent (*it, district, precinct))
	  ++numBallots;
      }
      assert (it->lastChoice >= it->firstChoice);
      const unsigned long M (it->lastChoice - it->firstChoice + 1);
      const unsigned long votesPerTriangle ((M*M+M)/2);
      it->triangleMultiplier = numBallots * it->N / votesPerTriangle;

      // In N-of-M, no choice can possibly receive more votes than the
      // number of ballots.
      if (it->logic == TGParser::Contest::NofM &&
          it->triangleMultiplier * M > numBallots)
        it->triangleMultiplier = numBallots / M;
    }
}


void generateVoterInput (ContestVector &contests,
			 const std::vector<unsigned long> &districtMap,
                         const std::vector<unsigned long> &styleMap) {
  figureMultipliers (contests, districtMap, styleMap);
  std::cout << "copy VoterInput from stdin CSV;\n";
  for (unsigned long ballot (1); ballot <= election.ballots; ++ballot) {
    const unsigned long precinct (styleMap[ballot]);
    const unsigned long district (districtMap[precinct]);
    for (ContestVector::iterator it (contests.begin());
	 it != contests.end();
         ++it)
      if (pertinent (*it, district, precinct))
        doContest (ballot, *it);
  }
  std::cout << "\\.\n\n";
}


int main (int argc, char **argv) {

  if (argc != 2)
    usage();

  yyin = fopen (argv[1], "r");
  if (!yyin) {
    perror (argv[1]);
    usage();
  }
  yyparse();

  unsigned long contestCount(1), contestSerial(1), choiceSerial(1);
  ContestVector contests;
  for (TGParser::ContestVector::iterator it (election.contests.begin());
       it != election.contests.end();
       ++it) {
    makeContest (contests, *it, contestSerial, choiceSerial);
    ++contestCount;
  }

  std::cout << "\\i Infrastructure-TestHeader.sql\n\n";
  std::cout << "\\echo 'TestGenerator-generated test.  Test spec follows:'\n";
  std::cout << "\\echo\n";
  {
    std::ifstream in (argv[1]);
    if (!in) {
      perror (argv[1]); // Unclear whether ifstream is required to set errno.
      usage();
    }
    std::string linebuf;
    while (getline (in, linebuf))
      std::cout << "\\echo '  " << linebuf << "'\n";
  }
  std::cout << "\n\\i Infrastructure-VoteSchema.sql\n\n";

  std::vector<std::string> reportingContexts;
  makeReportingContexts (reportingContexts);

  // districtMap maps from precinct ID to district ID.
  std::vector<unsigned long> districtMap;
  makePrecincts (districtMap);

  // styleMap maps from ballot ID to precinct ID (which also identifies the
  // ballot style).
  std::vector<unsigned long> styleMap;
  makeBallots (styleMap);

  generateReportingContexts (reportingContexts);
  generateContests (contests);
  generateChoices (contests);
  generateBallotStyles (contests, districtMap);
  generateBallots (styleMap);
  generateVoterInput (contests, districtMap, styleMap);

  std::cout << "\\i Infrastructure-IntegrityChecks.sql\n";

  std::cout << "\\! ReportGenerator/ReportGenerator";
  for (std::vector<std::string>::const_iterator it (reportingContexts.begin());
       it != reportingContexts.end();
       ++it)
    std::cout << " \"" << *it << '"';

  std::cout << "\n\\i Infrastructure-TestFooter.sql\n";

  return 0;
}
