// compute_error.cc - Don Yang (uguu.org)
//
// Compute error values for a particular block size, writes computed error
// values to stdout in CSV format.  This is used to tune crack parameters.
//
// Output columns:
// 1. block size
// 2. offset
// 3. stride
// 4. unshifted error
// 5. unshifted rank
// 6. minimum error
// 7. maximum error
// 8. average error
// 9. standard deviation
// 10. margin0 (raw)
// 11. margin1 (raw)
// 12. margin0 (standard deviations)
// 13. margin1 (standard deviations)
//
// 2015-06-20

#include<math.h>
#include<stdio.h>
#include<algorithm>
#include<array>
#include<string>

namespace {

using namespace std;

// Expected character distributions from "Alice In Wonderland"
// http://www.gutenberg.org/ebooks/11
static const int kCharCount = 26;
static const double kCharDistribution[kCharCount] =
{
   0.079698184390474,      // a
   0.0141963915472115,     // b
   0.0243924253388514,     // c
   0.0444673913927262,     // d
   0.125173796030539,      // e
   0.019367585719048,      // f
   0.0239289692574133,     // g
   0.064143947832733,      // h
   0.0701932693167682,     // i
   0.0019107399848767,     // j
   0.0104887428957061,     // k
   0.042369642813585,      // l
   0.0200587044369822,     // m
   0.065461138801031,      // n
   0.0770556716454317,     // o
   0.0160014310222865,     // p
   0.00178877785818244,    // q
   0.053744643829936,      // r
   0.0590947157875908,     // s
   0.0991958630446625,     // t
   0.032344355999317,      // u
   0.00782996853377131,    // v
   0.0240021465334298,     // w
   0.00143102228654595,    // x
   0.0210100090251974,     // y
   0.000650464675702705    // z
};

struct ErrorStats
{
   // Error value between unshifted input text and expected
   // distribution.
   double unshifted_error;

   // Rank of unshifted error among all sorted errors.  Usually 0.
   // This is used to measure how many candidate shift positions need
   // to be rejected to arrive at the correct key.
   int unshifted_error_rank;

   // Minimum error value between shifted input text and expected
   // distribution.  Normally this should be same as unshifted_error.
   double minimum_error;

   // Maximum error value between shifted input text and expected
   // distribution.
   double maximum_error;

   // Margin between minimum_error and second smallest error value.
   double margin0;

   // Margin between second smallest error value and third smallest
   // error value.
   double margin1;

   // Average of all error values.
   double average_error;

   // Standard deviation of all shifted error values.
   double standard_deviation;
};

// Compute stats for a block of text
static ErrorStats ComputeStats(
   const string &input, int offset, int block_size, int stride)
{
   ErrorStats stats;

   // Get character frequencies
   array<int, kCharCount> freq;
   freq.fill(0);
   for(int i = 0; i < block_size; i += stride)
      freq[input[offset + i]]++;

   // Convert frequencies to probabilities
   const double char_count = static_cast<double>(block_size / stride);
   array<double, kCharCount> prob;
   for(int i = 0; i < kCharCount; i++)
      prob[i] = static_cast<double>(freq[i]) / char_count;

   // Compute error values for each shift amounts
   vector<double> shifted_error;
   shifted_error.reserve(kCharCount);
   for(int shift = 0; shift < kCharCount; shift++)
   {
      double error = 0.0;
      for(int i = 0; i < kCharCount; i++)
      {
         const double delta = prob[(i + shift) % 26] - kCharDistribution[i];
         error += delta * delta;
      }
      shifted_error.push_back(error);
   }
   stats.unshifted_error = shifted_error.front();

   sort(shifted_error.begin(), shifted_error.end());
   stats.minimum_error = shifted_error.front();
   stats.maximum_error = shifted_error.back();

   // Compute derived stats
   for(int i = 0; i < kCharCount; i++)
   {
      if( shifted_error[i] == stats.unshifted_error )
      {
         stats.unshifted_error_rank = i;
         break;
      }
   }
   stats.margin0 = shifted_error[1] - shifted_error[0];
   stats.margin1 = shifted_error[2] - shifted_error[1];
   stats.average_error = 0.0;
   for(double e : shifted_error)
      stats.average_error += e;
   stats.average_error /= kCharCount;

   stats.standard_deviation = 0.0;
   for(double e : shifted_error)
   {
      const double delta = e - stats.average_error;
      stats.standard_deviation += delta * delta;
   }
   stats.standard_deviation = sqrt(stats.standard_deviation / kCharCount);

   return stats;
}

// Generate stats for a particular key size
static void ComputeStatsForKeySize(
   const string &input, int offset, int block_size, int key_size)
{
   for(int i = 0; i < key_size; i++)
   {
      const ErrorStats stats =
         ComputeStats(input, offset + i, block_size - i, key_size);
      printf("%d,%d,%d,%f,%d,%f,%f,%f,%f,%f,%f,",
             block_size,
             offset + i,
             key_size,
             stats.unshifted_error,
             stats.unshifted_error_rank,
             stats.minimum_error,
             stats.maximum_error,
             stats.average_error,
             stats.standard_deviation,
             stats.margin0,
             stats.margin1);
      if( stats.standard_deviation < 1e-6 )
      {
         puts(",");
      }
      else
      {
         printf("%f,%f\n",
                stats.margin0 / stats.standard_deviation,
                stats.margin1 / stats.standard_deviation);
      }
   }
}

// Generate error stats for various key sizes and offsets in this file
static void SampleStats(const string &input, int block_size)
{
   static const int kSamplesPerKey = 100;

   for(int key_size = 1; key_size <= 32; key_size++)
   {
      for(int i = 0; i < kSamplesPerKey; i++)
      {
         ComputeStatsForKeySize(
            input,
            i * (input.size() - block_size) / (kSamplesPerKey - 1),
            block_size,
            key_size);
      }
   }
}

// Load file to memory and convert alpha characters to binary
static string LoadFile(const char *filename)
{
   FILE *file = fopen(filename, "rb");

   if( file == nullptr )
   {
      printf("Error reading %s\n", filename);
      return "";
   }

   string data;
   char buffer[4096];
   for(int read_size = sizeof(buffer);
       read_size == sizeof(buffer) && !feof(file);)
   {
      read_size = fread(buffer, 1, sizeof(buffer), file);
      for(int i = 0; i < read_size; i++)
      {
         if( buffer[i] >= 'a' && buffer[i] <= 'z' )
            data.push_back(buffer[i] - 'a');
         else if( buffer[i] >= 'A' && buffer[i] <= 'Z' )
            data.push_back(buffer[i] - 'A');
      }
   }
   fclose(file);

   return data;
}

}  // namespace

int main(int argc, char **argv)
{
   int block_size;
   if( argc != 3 ||
       sscanf(argv[2], "%d", &block_size) != 1 ||
       block_size < 100 )
   {
      return printf("%s <file.txt> <block_size>\n", *argv);
   }

   const string data = LoadFile(argv[1]);
   if( static_cast<int>(data.size()) < block_size )
   {
      printf("Insufficient alpha characters in %s (%zu < %d)\n",
             argv[1], data.size(), block_size);
      return 1;
   }

   SampleStats(data, block_size);
   return 0;
}