
#include "polybius.hpp"
#include <iostream>
#include <cstdlib>

using namespace std;

PolybiusSquare::PolybiusSquare( enum MODE mode, size_t dimension )
    :
    m_Mode          ( mode ),
    m_KeyX          ( "NO KEY" ),
    m_KeyY          ( "NO KEY" ),
    m_Verbose       ( false ),
    m_PRNGReshuffle ( false ),
    m_PRNGIndex     ( 0 )
{
    m_Dimension = dimension>=6? dimension:6;

    m_Data = new char[ m_Dimension * m_Dimension ];
    m_KeyX.resize( m_Dimension );
    m_KeyY.resize( m_Dimension );

    switch( m_Mode )
    {
        case MODE_RNG:
        {
            size_t index = 0;
            int    entry = 0;

            for( int n=0; n<(int)(m_Dimension*m_Dimension); n++)
            {
                m_Data[n]='0'+(entry%10);
                entry++;
            }
            break;
        }
        case MODE_ALPHANUM:
        {
            size_t index = 0;
            char   entry = '0';
            int    cnt   = 0;

            // fill in numbers
            while( entry <= '9' )
            {
                m_Data[index] = entry;
                index++;
                entry++;
            }

            entry = m_Dimension == 6? 'A':'['; // start with " " if bigger than
                                               // 6x6 otherwise start with "A"

            while( index < ( m_Dimension * m_Dimension ) )
            {
                // A, I, O, U, are filled in twice E is filled in
                // 4 times.

		// Remark: If the Polybius square is very big, it might
		//         be advisable to consider T and H, too, or
		//         using an even better distribution.
                if( entry == 'A' or
                    entry == 'I' or
                    entry == 'O' or
                    entry == 'U' )
                {
                    cnt = 2;
                }
                else if ( entry == 'E' )
                {
                    cnt = 4;
                }
                else
                {
                    cnt = 1;
                }

                // If the Polybius square is very small, omit repeated vocals.
                // This is very bad but there is not enough space left.
                if( m_Dimension == 6 )
                {
                    cnt=1;
                }

                while( cnt > 0 )
                {
                    if( index < ( m_Dimension * m_Dimension ) )
                    {
                        if( entry == '[' )
                        {
                            m_Data[ index ] = ' ';
                        }
                        else
                        {
                            m_Data[ index ] = entry;
                        }
                        index++;
                    }
                    cnt--;
                }
                entry++;
                if( entry == '\\' )
                {
                    entry = 'A';
                }
            }

            break;
        }
    }
}

PolybiusSquare::~PolybiusSquare()
{
    delete [] m_Data;
    m_Data = 0;
    m_Dimension = 0;
}

void PolybiusSquare::setVerbose( bool flag )
{
    m_Verbose = flag;
}

void PolybiusSquare::setPRNGReshuffle()
{
    m_PRNGReshuffle = true;
}

void PolybiusSquare::print()
{
    cout << endl;

    cout << endl << "      ";
    for( int x=0; x<(int)m_Dimension; ++x )
    {
        cout << (char)((x<=9)? x+'0':x-10+'A');
    }

    cout << endl << "      ";
    for( int x=0; x<(int)m_Dimension; ++x )
    {
        cout << '-';
    }

    cout << endl << "      ";
    for( int x=0; x<(int)m_Dimension; ++x )
    {
        cout << m_KeyX[ x ];
    }

    cout << endl << "      ";
    for( int x=0; x<(int)m_Dimension; ++x )
    {
        cout << '-';
    }

    cout << endl;
    for( int y=0; y<(int)m_Dimension; ++y )
    {
        cout << "  "
             << (char)((y<=9)? y+'0':y-10+'A')
             << "|"
             << m_KeyY[ y ]
             << "|";

        for( int x=0; x<(int)m_Dimension; ++x )
        {
            cout << m_Data[ x + y * m_Dimension ];
        }
        cout << endl;
    }
}

std::string PolybiusSquare::encode( char c )
{
    int readIndex = 0;
    std::string coordinate;

    c = toupper( c );

    if( not (
            ( c >= '0' and c <= '9' ) or
            ( c >= 'A' and c <= 'Z' ) or
            ( c == ' ' )
            ) )
    {
        return( "" );
    }

    if( ( m_Dimension == 6 ) and ( c == ' ' ) )
    {
        return( "" );
    }

    // this is not very stochastic... but would a human user be? Fortunately we
    // don't need strong randomness for this to work well.
    readIndex = rand()%(m_Dimension*m_Dimension);

    int foundAt=-1;
    for( int n=0; n<(int)(m_Dimension*m_Dimension); ++n )
    {
        int idx = ( readIndex + n )%(m_Dimension*m_Dimension);
        if( m_Data[idx] == c )
        {
            foundAt = idx;
            break;
        }
    }
    if( foundAt < 0 )
    {
        cout << "Warning: character '" << c << "' not found." << endl;
    }
    int y = (int)foundAt / (int)m_Dimension;
    int x = (int)foundAt % (int)m_Dimension;

    coordinate.push_back( x<10? (char)('0'+x):(char)('A'+x-10) );
    coordinate.push_back( y<10? (char)('0'+y):(char)('A'+y-10) );

    return( coordinate );
}

std::string PolybiusSquare::encode( std::string s )
{
    std::string out("");

    for( int n=0; n<(int)s.size(); ++n )
    {
        out += encode( s[n] );
    }

    return( out );
}

std::string PolybiusSquare::decode( std::string s )
{
    std::string out("");

    for( int n=0; n<(int)s.size(); n += 2 )
    {
        out += m_Data[ (s[n+0]-'0')+(s[n+1]-'0')*m_Dimension ];
    }

    return( out );
}

std::string PolybiusSquare::addPRNG( std::string s )
{
    std::string result("");

    for( int n=0; n<(int)s.size(); ++n )
    {
        int index = n%(m_Dimension*m_Dimension);

        int input = ( s[n] - '0' ) % 10;
        int prng  = ( m_Data[index] - '0' ) % 10;
        int out   = ( input + prng ) % 10;

        result += (char)(out+'0');

        if( (n%100)==99 && m_PRNGReshuffle )
        {
            transpose();
        }
    }
    return( result );
}

std::string PolybiusSquare::subPRNG( std::string s )
{
    std::string result("");

    for( int n=0; n<(int)s.size(); ++n )
    {
        int index = n%(m_Dimension*m_Dimension);

        int input = ( s[n] - '0' ) % 10;
        int prng  = ( m_Data[index] - '0' ) % 10;
        int out   = ( input + 10 - prng ) % 10;

        result += (char)(out+'0');

        if( (n%100)==99 && m_PRNGReshuffle )
        {
            transpose();
        }
    }
    return( result );
}

void PolybiusSquare::setKey(std::string key)
{
    // convert key to upper case
    for( auto & c: key )
    {
        c = toupper(c);
    }

    int index=0;
    for( int n=0; n<(int)m_Dimension; ++n )
    {
        m_KeyX[n] = key[index];

        index++;

        if( index >= (int)key.size() )
        {
            index = 0;
        }
    }
    for( int n=0; n<(int)m_Dimension; ++n )
    {
        m_KeyY[n] = key[index];

        index++;

        if( index >= (int)key.size() )
        {
            index = 0;
        }
    }
}

void PolybiusSquare::transposeY()
{
    for( int x=0; x<(int)m_Dimension; ++x )
    {
        int shift = (m_KeyX[x] - 'A' )%10;

        char* shiftedCol = new char[m_Dimension];

        for( int y=0; y<(int)m_Dimension; ++y )
        {
            int yt = (y+(int)m_Dimension-shift)%(int)m_Dimension;
            shiftedCol[y] = m_Data[x+(yt)*m_Dimension];
        }

        for( int y=0; y<(int)m_Dimension; ++y )
        {
            m_Data[x+(y)*m_Dimension] = shiftedCol[y];
        }

        delete [] shiftedCol;
    }
}

void PolybiusSquare::transposeX()
{
    for( int y=0; y<(int)m_Dimension; ++y )
    {
        int shift = (m_KeyY[y] - 'A')%10;

        char* shiftedRow = new char[m_Dimension];

        for( int x=0; x<(int)m_Dimension; ++x )
        {
            int xt = (x+(int)m_Dimension-shift)%(int)m_Dimension;
            shiftedRow[x] = m_Data[xt+y*m_Dimension];
        }

        for( int x=0; x<(int)m_Dimension; ++x )
        {
            m_Data[x+y*m_Dimension] = shiftedRow[x];
        }

        delete [] shiftedRow;
    }
}

void PolybiusSquare::transpose()
{
    if( m_Verbose )
    {
        cout << endl << endl;
        cout << "=====================================================" << endl;
        cout << "Initial Polybius square: " << endl;
        print();
    }

    transposeX();

    if( m_Verbose )
    {
        cout << endl;
        cout << "Polybius square after rotating rows (transposeX): "
             << endl;
        print();
    }

    transposeY();

    if( m_Verbose )
    {
        cout << endl;
        cout << "Polybius square after rotating columns (transposeY): "
             << endl;
        print();
    }
}

