#include "stdafx.h"
#include "Square.h"
#include "GaussianInteger.h"
#include "time.h"

using namespace std;
using namespace NTL;

int MIN_PRIME;
int MAX_PRIME;
int MAX_POWER;
int MAX_FACTORS;
int MIN_NUMBER_OF_SQUARES;

int* prime;
int* primesOneModFour;
int spot;
int primesOneModFourCount;

GaussianInteger*** sumOfSquaresRepresentationTable = 0;

fstream outFile;

int _tmain(int argc, _TCHAR* argv[]) // MAX_PRIME MAX_POWER MAX_FACTORS MIN_NUMBER_OF_SQUARES
{
MIN_PRIME = 0;
MAX_PRIME = 100000;
MAX_POWER = 2*1 + 1;
MAX_FACTORS = 5;
MIN_NUMBER_OF_SQUARES = 6;
readArguments( argc, argv );

prime = new int[MAX_PRIME];
primesOneModFour = new int[MAX_PRIME];

srand( (unsigned) time(NULL) );
mt_init();

ostringstream stringStream;
stringStream << "log" << MIN_PRIME << "-" << MAX_PRIME << "-" << MAX_POWER << "-" << MAX_FACTORS << "-" << MIN_NUMBER_OF_SQUARES << ".txt";
outFile.open( stringStream.str().c_str(), fstream::in | fstream::out | fstream::ate );

int i;
int j;

bool bPreviousWorkSucked = true;
Factor* factorList = new Factor[MAX_FACTORS];
factorList[0].power = 2 * ( MAX_POWER / 2 );
for(i = 1; i < MAX_FACTORS; i++ )
factorList[i].power = 2 * ( MAX_POWER / 2 );

outFile.seekg( 0, fstream::end );
int length = outFile.tellg();
if( length > 0 )
{
cout << "Some work has already been done on this range, processing log file...\n";
outFile << "\nSome work has already been done on this range, processing log file...\n";

outFile.seekg( -2000, fstream::end );
string line[100];
int count = 0;
while( getline( outFile, line[count] ) )
{
// cout << "line(" << count << "): " << line[count] << "\n";
if( line[count].find( ":" ) != string::npos )
count++;
}


if( count >= 1 )
{

string temp = line[count - 1];
int index;
int start = 0;
for( i = 0; i < MAX_FACTORS; i++ )
{
index = temp.find( "^", start );
if( index == string::npos )
break;
stringstream tempStringStream;
tempStringStream << temp.substr( start, index - start );
tempStringStream >> factorList[i].prime;
// cout << "\nstreamed this in: " << factorList[i].prime << " from this: " << tempStringStream.str() << "\n";
start = index + 1;

start = temp.find( ".", start ) + 1;
if( index == string::npos )
break;
}
if( i == MAX_FACTORS )
bPreviousWorkSucked = false;
}
}

outFile.clear();
outFile.close();

outFile.open( stringStream.str().c_str(), fstream::out | fstream::app );

cout << "Generating all primes less than " << MAX_PRIME << ".\n";
outFile << "Generating all primes less than " << MAX_PRIME << ".\n";
generatePrimes();
cout << "Precomputing sum of squares decompositions for prime powers up to exponent " << MAX_POWER - 1 << ".\n";
outFile << "Precomputing sum of squares decompositions for prime powers up to exponent " << MAX_POWER - 1 << ".\n";
setupSumOfSquaresTable();

if( !bPreviousWorkSucked )
{
for( i = 0; i < MAX_FACTORS; i++ )
{
for( j = 0; j < primesOneModFourCount; j++ )
if( factorList[i].prime == primesOneModFour[j] )
{
factorList[i].prime = j;
break;
}
if( j == primesOneModFourCount )
{
cout << "\nfailed to find prime: " << factorList[i].prime << "\n";
bPreviousWorkSucked = true;
break;
}
}
}

if( bPreviousWorkSucked )
{
factorList[0].prime = 0;
for(i = 1; i < MAX_FACTORS; i++ )
factorList[i].prime = factorList[i - 1].prime + 1;
}

int numberOfFactors = MAX_FACTORS;
ZZ M;
ZZ tempA;
ZZ tempB;
bool done = false;
bool carrying;
GaussianInteger* decompositionList = 0;
bool foundSomething = false;
while( !done )
{
if( !done )
{
for( i = 0; i < numberOfFactors - 1; i++ )
{
cout << primesOneModFour[factorList[i].prime] << "^" << factorList[i].power << ".";
outFile << primesOneModFour[factorList[i].prime] << "^" << factorList[i].power << ".";
}
cout << primesOneModFour[factorList[i].prime] << "^" << factorList[i].power << ":\n";
outFile << primesOneModFour[factorList[i].prime] << "^" << factorList[i].power << ":\n";

cout << "Computing all sum of squares decompositions... ";
outFile << "Computing all sum of squares decompositions... ";

int numberOfDecompositions = getAllDecompositions( factorList, numberOfFactors, decompositionList );

if( spot != numberOfDecompositions )
{
cout << "I mispredicted the number of sum of squares decompositions!\n";
outFile << "I mispredicted the number of sum of squares decompositions!\n";
break;
}

cout << "found " << numberOfDecompositions << " decompositions.\n";
outFile << "found " << numberOfDecompositions << " decompositions.\n";

ZZ middle;
middle = numericallValueOfFactorList( factorList, numberOfFactors );
M = 3 * middle;

cout << "Building all possible magic squares...\n";
outFile << "Building all possible magic squares...\n";
ZZ square[9]; // 0 1 2
// 3 4 5
// 6 7 8
square[4] = middle;

for( i = 0; i < numberOfDecompositions; i++ )
{
for(j = i + 1; j < numberOfDecompositions; j++ )
{
// skew
square[6] = decompositionList[i].a;
square[2] = decompositionList[i].b;
square[7] = decompositionList[j].a;
square[1] = decompositionList[j].b;
square[0] = M - ( square[1] + square[2] );
square[8] = M - ( square[0] + square[4] );
square[5] = M - ( square[8] + square[2] );
square[3] = M - ( square[5] + square[4] );
if( isMagic( square ) )
{
int nSquares = numberOfSquares( square );
if( nSquares >= MIN_NUMBER_OF_SQUARES )
{
foundSomething = true;
cout << "found something (" << nSquares << "):\n";
outFile << "found something (" << nSquares << "):\n";
printSquare( square );
}
}

square[6] = decompositionList[i].a;
square[2] = decompositionList[i].b;
square[7] = decompositionList[j].b;
square[1] = decompositionList[j].a;
square[0] = M - ( square[1] + square[2] );
square[8] = M - ( square[0] + square[4] );
square[5] = M - ( square[8] + square[2] );
square[3] = M - ( square[5] + square[4] );
if( isMagic( square ) )
{
int nSquares = numberOfSquares( square );
if( nSquares >= MIN_NUMBER_OF_SQUARES )
{
foundSomething = true;
cout << "found something[b] (" << nSquares << "):\n";
outFile << "found something[b] (" << nSquares << "):\n";
printSquare( square );
}
}

square[6] = decompositionList[j].a;
square[2] = decompositionList[j].b;
square[7] = decompositionList[i].a;
square[1] = decompositionList[i].b;
square[0] = M - ( square[1] + square[2] );
square[8] = M - ( square[0] + square[4] );
square[5] = M - ( square[8] + square[2] );
square[3] = M - ( square[5] + square[4] );
if( isMagic( square ) )
{
int nSquares = numberOfSquares( square );
if( nSquares >= MIN_NUMBER_OF_SQUARES )
{
foundSomething = true;
cout << "found something[c] (" << nSquares << "):\n";
outFile << "found something[c] (" << nSquares << "):\n";
printSquare( square );
}
}

square[6] = decompositionList[j].a;
square[2] = decompositionList[j].b;
square[7] = decompositionList[i].b;
square[1] = decompositionList[i].a;
square[0] = M - ( square[1] + square[2] );
square[8] = M - ( square[0] + square[4] );
square[5] = M - ( square[8] + square[2] );
square[3] = M - ( square[5] + square[4] );
if( isMagic( square ) )
{
int nSquares = numberOfSquares( square );
if( nSquares >= MIN_NUMBER_OF_SQUARES )
{
foundSomething = true;
cout << "found something[d] (" << nSquares << "):\n";
outFile << "found something[d] (" << nSquares << "):\n";
printSquare( square );
}
}

// 0 1 2
// 3 4 5
// 6 7 8

// plus
square[3] = decompositionList[i].a;
square[5] = decompositionList[i].b;
square[7] = decompositionList[j].a;
square[1] = decompositionList[j].b;
square[0] = ( square[7] + square[5] ) / 2;
square[2] = ( square[3] + square[7] ) / 2;
square[6] = ( square[1] + square[5] ) / 2;
square[8] = ( square[3] + square[1] ) / 2;
if( isMagic( square ) )
{
int nSquares = numberOfSquares( square );
if( nSquares >= MIN_NUMBER_OF_SQUARES )
{
foundSomething = true;
cout << "found something[e] (" << nSquares << "):\n";
outFile << "found something[e] (" << nSquares << "):\n";
printSquare( square );
}
}

// X
square[0] = decompositionList[i].a;
square[8] = decompositionList[i].b;
square[6] = decompositionList[j].a;
square[2] = decompositionList[j].b;
square[1] = M - ( square[0] + square[2] );
square[3] = M - ( square[0] + square[6] );
square[5] = M - ( square[2] + square[8] );
square[7] = M - ( square[6] + square[8] );
if( isMagic( square ) )
{
int nSquares = numberOfSquares( square );
if( nSquares >= MIN_NUMBER_OF_SQUARES )
{
foundSomething = true;
cout << "found something[f] (" << nSquares << "):\n";
outFile << "found something[f] (" << nSquares << "):\n";
printSquare( square );
}
}
}
}
cout << "\n";
outFile << endl;
outFile.flush();
}

// Acquire the next central square to try.
j = 1;
do
{
carrying = false;
factorList[MAX_FACTORS - j].prime++;
if( factorList[MAX_FACTORS - j].prime >= primesOneModFourCount - j )
{
j++;
if( j <= MAX_FACTORS )
carrying = true;
else
done = true;
}
}
while( carrying );

if( !done )
{
for( i = MAX_FACTORS - j + 1; i < MAX_FACTORS; i++ )
{
factorList[i].prime = factorList[i - 1].prime + 1;
}
}
}
if( foundSomething )
{
cout << "A viable magic square with the requisite number of squares was found.\n";
outFile << "A viable magic square with the requisite number of squares was found.\n";
}
outFile.close();
if( decompositionList )
delete []decompositionList, decompositionList = 0;
if( factorList )
delete []factorList, factorList = 0;
if( prime )
delete []prime, prime = 0;
if( primesOneModFour )
delete []primesOneModFour, primesOneModFour = 0;
return 0;
}

void readArguments(int argc, _TCHAR* argv[])
{
if( argc >= 2 )
{
wistringstream tempStream( argv[1] );
tempStream >> MIN_PRIME;
if( argc >= 3 )
{
wistringstream tempStream( argv[2] );
tempStream >> MAX_PRIME;
if( argc >= 4 )
{
wistringstream tempStream( argv[3] );
tempStream >> MAX_POWER;
MAX_POWER = 2 * (MAX_POWER / 2) + 1;
if( argc >= 5 )
{
wistringstream tempStream( argv[4] );
tempStream >> MAX_FACTORS;
if( argc >= 6 )
{
wistringstream tempStream( argv[5] );
tempStream >> MIN_NUMBER_OF_SQUARES;
}
}
}
}
}
}

bool isMagic( ZZ* square )
{
// 0 1 2
// 3 4 5
// 6 7 8
int i;
int j;
for( i = 0; i < 9; i++ )
for( j = i + 1; j < 9; j++ )
if( square[i] == square[j] )
return false;

ZZ M = 3 * square[4];
return ( square[0] + square[1] + square[2] == M ) &&
( square[3] + square[4] + square[5] == M ) &&
( square[6] + square[7] + square[8] == M ) &&
( square[0] + square[3] + square[6] == M ) &&
( square[1] + square[4] + square[7] == M ) &&
( square[2] + square[5] + square[8] == M ) &&
( square[0] + square[4] + square[8] == M ) &&
( square[6] + square[4] + square[2] == M );

}

void printSquare( ZZ* square )
{
int i;
ZZ root;
bool isSquare;
cout << "\n";
outFile << "\n";
for( i = 0; i < 9; i++ )
{
isSquare = false;
if( square[i] >= 0 )
{
root = SqrRoot( square[i] );
if( root * root == square[i] )
isSquare = true;
}
if( isSquare )
{
cout << root << "^2 ";
outFile << root << "^2 ";
}
else
{
cout << square[i] << " ";
outFile << square[i] << " ";
}
if( ( i + 1 ) % 3 == 0 )
{
cout << "\n";
outFile << "\n";
}
}
cout << "\n";
outFile << "\n";
}

int numberOfSquares( ZZ* square )
{
int numberOfSquaresFound = 0;
int i;
ZZ root;
for( i = 0; i < 9; i++ )
{
if( square[i] >= 0 )
{
root = SqrRoot( square[i] );
if( root * root == square[i] )
numberOfSquaresFound++;
}
}
return numberOfSquaresFound;
}

void decomposePrime(const ZZ& p, ZZ& a, ZZ& b)
{
ZZ left;
ZZ right;
ZZ temp;

if ( p % 8 == 5 )
right = 2;
else
{
right = to_ZZ( 3 ) % p;
while ( Jacobi( right, p ) == 1 )
{
temp = nextPrime( to_int( right ) );
if( temp == -1 )
{
cout << "\nRAN OUT OF PRIMES!\n";
return;
}
right = temp % p;
}
}

right = PowerMod( right, ( p - 1 ) / 4, p );

if ( right < p / 2 )
right = p - right;
left = p;


while ( right * right > p )
{
temp = left;
left = right;
right = temp % right;
}
a = SqrRoot( p - right * right );
b = right;
}

ZZ nextPrime( int currentPrime )
{
int i;
for( i = currentPrime + 1; i < MAX_PRIME; i++ )
{
if ( prime[i] )
return to_ZZ( i );
}
return to_ZZ( -1 );
}

void generatePrimes()
{
int i, j;

prime[0] = 0;
prime[1] = 0;
for( i = 2; i < MAX_PRIME; i++ )
prime[i] = 1;
j = 0;
while( j * j <= MAX_PRIME )
{
while( prime[j] == 0 && j * j <= MAX_PRIME )
j++;
for( i = j + 1; i < MAX_PRIME; i++ )
{
if ( i % j == 0 )
prime[i] = 0;
}
j++;
}

primesOneModFourCount = 0;

for( i = 1; i < MAX_PRIME; i += 4 )
{
if( prime[i] && i >= MIN_PRIME)
{
primesOneModFour[primesOneModFourCount] = i;
primesOneModFourCount++;
}
}

}

int getAllDecompositions( Factor* factorList, int numberOfFactors, GaussianInteger*& decompositionList )
{
int i;
int numberOfDecompositions = 1;

for( i = 0; i < numberOfFactors; i++ )
numberOfDecompositions *= factorList[i].power + 1;

numberOfDecompositions /= 2;

if( decompositionList )
delete []decompositionList, decompositionList = 0;
decompositionList = new GaussianInteger[numberOfDecompositions];

GaussianInteger x( 1, 0 );
spot = 0;
recurse( factorList, numberOfFactors, decompositionList, x, 0 );

return numberOfDecompositions;
}

void recurse( Factor* factorList, int numberOfFactors, GaussianInteger* decompositionList,
const GaussianInteger& x, int currentFactor )
{
int i;
if( currentFactor == numberOfFactors ) // no factors left
{
GaussianInteger r( 1, 1 );
GaussianInteger ostensiblyNewDecompostion = x * r;

ostensiblyNewDecompostion.a = ostensiblyNewDecompostion.a * ostensiblyNewDecompostion.a;
ostensiblyNewDecompostion.b = ostensiblyNewDecompostion.b * ostensiblyNewDecompostion.b;

// make sure it really is a new and viable decomposition
if( ostensiblyNewDecompostion.a == ostensiblyNewDecompostion.b )
return;
for( i = 0; i < spot; i++ )
{
if( ostensiblyNewDecompostion.a == decompositionList[i].a && ostensiblyNewDecompostion.b == decompositionList[i].b ||
ostensiblyNewDecompostion.b == decompositionList[i].a && ostensiblyNewDecompostion.a == decompositionList[i].b )
return;
}

decompositionList[spot] = ostensiblyNewDecompostion;
spot++;
}
else
{
GaussianInteger y;
GaussianInteger z;
for( i = 0; i <= factorList[currentFactor].power; i++)
{
y = x * sumOfSquaresRepresentationTable[factorList[currentFactor].prime][factorList[currentFactor].power][i];
recurse( factorList, numberOfFactors, decompositionList, y, currentFactor + 1 );
}
}
}

void setupSumOfSquaresTable()
{
int i;
int j;
int k;
sumOfSquaresRepresentationTable = new GaussianInteger**[primesOneModFourCount];
for( i = 0; i < primesOneModFourCount; i++ )
{
sumOfSquaresRepresentationTable[i] = new GaussianInteger*[MAX_POWER];
for( j = 0; j < MAX_POWER; j++ )
sumOfSquaresRepresentationTable[i][j] = new GaussianInteger[j + 1];
}

GaussianInteger x;
GaussianInteger y;
for( i = 0; i < primesOneModFourCount; i++ )
{
decomposePrime( to_ZZ( primesOneModFour[i] ), x.a, x.b );
y = x.conjugate();
for( j = 1; j < MAX_POWER; j++ )
for( k = 0; k <= j; k++ )
sumOfSquaresRepresentationTable[i][j][k] = x.toThe( k ) * y.toThe( j - k );
}
}

void killSumOfSquaresTable()
{
int i;
int j;
for( i = 0; i < primesOneModFourCount; i++ )
{
for( j = 0; j < MAX_POWER; j++ )
{
if( sumOfSquaresRepresentationTable[i][j] )
delete []sumOfSquaresRepresentationTable[i][j],sumOfSquaresRepresentationTable[i][j] = 0;
}
if( sumOfSquaresRepresentationTable[i] )
delete sumOfSquaresRepresentationTable[i], sumOfSquaresRepresentationTable[i] = 0;
}
if( sumOfSquaresRepresentationTable )
delete sumOfSquaresRepresentationTable, sumOfSquaresRepresentationTable = 0;
}