#include <iostream>
#include <vector>
#include "pdist.h"

using namespace std;



inline vector<double> operator+(const vector<double> &l, const vector<double> &r)
{
    vector<double> res(l.size());

    for(int i = 0; i < l.size(); i++)
        res[i] = l[i] + r[i];

    return res;
}

inline vector<double> operator*(const double &l, const vector<double> &r)
{
    vector<double> res(r.size());

    for(int i = 0; i < r.size(); i++)
        res[i] = l * r[i];

    return res;
}

template <class T>
inline vector<T> &operator<<(vector<T> &l, const T &r)
{
    for(int i = 0; i < l.size(); i++)
        l[i] = r;

    return l;
}

inline void ComputeSwap(vector< vector<double> > &PS, int i, vector<double> &Pj)
{
    vector< vector<double> > PStminus1 = PS;

    // Compute probability dist of a swap into i
    
    // Matrix product (which?)
    // row vector . Matrix = row vector (PS[i] = Pj.PS)
    PS[i] << 0.0;

    // P(S[i] = s) 
    //  = P(J = k)*P(PStminus1[k] == s)
    for(int k = 0; k < Pj.size(); k++)
        PS[i] = PS[i] + (Pj[k] * PStminus1[k]);

    // Now compute probabilities at PS[j]
    // Pj and S have the same size, and S is always square (1<<n)
    for(int s = 0; s < Pj.size(); s++)
    {
        // P(S[s] == k) = 
        //  P(J = s)*P(Stminus1[i] = k) + P(J != s)*P(S[s] = k)

        // Or, in english, the probability that k is in an index of s is the
        // probability that j caused a swap into s and Sprev[i] was k, plus
        // the probability that J didn't cause a swap into s times the 
        // probability that S[s] already was k (after the swap for i)

        PS[s] = Pj[s]*PStminus1[i] + ((1-Pj[s])*PS[s]);
    }
}

void checkNormalization(vector<double> &n)
{
    double norm = 0;

    for(int i = 0; i < n.size(); i++)
        norm += n[i];
  
    if(!FEQ(norm, 1.0))
        cerr << "Normalization not 1: " << norm << endl;
}

// Plan: Comute a 2D Array with S[i][j] cooresponding to P(S[i] == j) 
// from an 2D array with PK[i][j] corresponding to P(K[i] == j).
#ifdef DEBUG_VERIFY_ROUNDS
vector< vector<double> > ComputePS(vector< vector<double> > &PK, int n, 
        const vector<int> &checkj)
#else
vector< vector<double> > ComputePS(vector< vector<double> > &PK, int n)
#endif
{
    // RC4 KSA:
    //  S = <0, 1, 2..N-1>
    //  j = 0
    //  for i = 1..N-1
    //    j = j + S[i] + K[i % l]
    //    Swap(S[i], S[j])
    const int N = 1<<n;
    const int mask = N-1;
   
    vector<double> Pj(N, 0.0);
    vector<double> Pjtminus1(N, 0.0);
    vector< vector<double> > PS(N, Pj);

    // P(j == 0) = 1
    Pj[0] = 1;
   
#ifdef DEBUG_VERIFY_NORMS
    cerr << "K normalization\n";
    for(int i = 0; i < PK.size(); i++)
        checkNormalization(PK[i]);
#endif
    
    // P(S[i] == i) = 1
    for(int i = 0; i < N; i++)
        PS[i][i] = 1;

    for(int i = 0; i < N; i++)
    {
        Pjtminus1 = Pj;

        Pj << 0.0;
      
        // P(j_{t-1}+S[i]+K[i] = j mod n) 
        //  = Sum over jp,s P(j_{t-1} == jp)*P(S[i] == s)*P(K[i] == j-s-jp)
        for(int j = 0; j < N; j++)
            for(int s = 0; s < N; s++)
                for(int jp = 0; jp < N; jp++)
                    Pj[j] = Pj[j] + 
                        Pjtminus1[jp]*PS[i][s]*PK[i%PK.size()][(((j-s-jp)%N)+N)%N];
        
        ComputeSwap(PS, i, Pj);


#ifdef DEBUG_VERIFY_ROUNDS
        cerr << endl;
        if(!FEQ(Pj[checkj[i]], 1.0))
            cerr << "Difference in round " << i << " at index " << checkj[i] << endl;
        

        for(int j = 0; j < N; j++)
            if(FEQ(Pj[j], 1.0))
                cerr << "j = " << j << "; ";
       
        cerr << endl;
        

        checkNormalization(Pj);
#endif
        cerr << "Completed round " << i << endl;
    }

#ifdef DEBUG_VERIFY_NORMS
    cerr << "S normalization\n";
    for(int s = 0; s < N; s++)
        checkNormalization(PS[s]);
    cerr << "end S normalization\n";

    cerr << "j normalization\n";
    checkNormalization(Pj);
#endif

    return PS;
}



// TODO: TEST
//  1. Does this dist do RC4 if given probabilities of 1
//  2. Does this dist come up with results comparable to Roo95 and Wag95, 
//  3. How do these results compare to the Random-Shuffle dist?
//  4. What is the complexity of computing prob dists throughout RC4?
//     Doesn't it add a factor of at least 2^8? (Possibly 2^10 or 12 with a 
//     slow FPU).
//
//     Does Knudson's probability dist add this much overhead? Is there a 
//     shortcut? Perhaps... Moment Generating Functions for the convolution
//     No good.. N^2.. Can you do FFT of a MGF though? Saves only 5 bits..