//
// $Id: pMC_mult.cpp,v 1.2 2006/04/28 17:35:52 nielsen Exp $
//
// (C) Jens Erik Nielsen, University College Dublin 2005
//
#include "pMC_mult.h"
#include "stdio.h"
//
//
//

void MC::reformat_arrays() {
    //
    // Reformat the matrix
    //
    //printf ("Reformatting arrays\n");
    _groups=static_cast<int>(_num_states.size());
    //printf ("Number of groups: %d\n",_groups);
    //for (int group=0;group<_groups;group++) {
    //  printf ("Num_states for group %d is %d\n",group,_num_states[group]); 
    //}

    //
    int count=0;
    for (int row=0;row<_groups;row++) {
        //printf ("Constructing matrix for group %d\n",row);
        vector<vector<vector<double > > > row_vals;
        for (int column=0;column<_groups;column++) {
            //printf ("Second level matrix for group %d\n",column);
            vector<vector<double > > column_vals;
            for (int group1_states=0;group1_states<_num_states[row];group1_states++) {
                vector<double> group1_s;
                for (int group2_states=0;group2_states<_num_states[column];group2_states++) {
                    group1_s.push_back(_lin_matrix[count]);

                    // printf ("While reformatting arrays: g1: %d, g2: %d, st1: %d st2: %d value: %5.3f\n",
                    //	  row,column,group1_states,group2_states,_lin_matrix[count]);
                    count=count+1;
                }
                column_vals.push_back(group1_s);
            }
            row_vals.push_back(column_vals);
        }
        _matrix.push_back(row_vals);
    }
    //
    // Reformat the intrinsic pKa array and the charged_state_array
    //
    count=0;
    for (int group=0;group<_groups;group++) {
        vector<double> these_intpKas;
        vector<int> these_charged_state;
        for (int t_state=0;t_state<_num_states[group];t_state++) {
            these_intpKas.push_back(_intpKas_lin[count]);
            these_charged_state.push_back(_charged_state_lin[count]);
            count++;
        }
        _intpKas.push_back(these_intpKas);
        _charged_state.push_back(these_charged_state);
    }
    //
    // Set natural log
    //
    lnten=log(10.0);
    return;
}

//
// ---------------------
//

vector<float> MC::calc_pKas(float pH_start,float pH_end, float pH_step) {
    //
    // Calculate pKa values for the system
    //
    // First get charges at all pH values
    //
    float max_pH=0.0;
    vector< vector<float> > charges;
    vector<float> calced_charges;
    for (float pH=pH_start;pH<pH_end;pH=pH+pH_step) {
        calced_charges = calc_charge(pH);
        charges.push_back(calced_charges);
        max_pH=pH;
    }
    printf("\n");
    // 
    // Now determine pKa values
    //
    int datapoints=11; // How many data points do we use for pKa determination?
    datapoints=(datapoints-1)/2;
    vector<float> pKas;
    for (int group=0;group<_groups;group++) {
        int count=0;
        float pKa=-999.9;
        float last_crg=charges[count][group];
        //
        for (float pH=pH_start;pH<pH_end;pH=pH+pH_step) {
            if ((pH-max_pH)>0.0) {
                continue;
            }
            float this_crg=charges[count][group];
            if (_acid_base[group]==1.0) {
                if (this_crg<=0.5 && last_crg>0.5) {
                    //
                    // Get ph,charge sets and calc pKa from those
                    //
                    vector<double> pHs_pKadet;
                    vector<float> charges_pKadet;
                    int count2=count-static_cast<int>(datapoints);
                    if (count2<0) {count2=0;}
                    for (double pH2=max(pH_start,pH-datapoints*pH_step);pH2<min(pH_end,pH+datapoints*pH_step);pH2=pH2+pH_step) {
                        pHs_pKadet.push_back(pH2);
                        charges_pKadet.push_back(charges[count2][group]);
                        count2=count2+1;
                    }
                    pKa=calc_pKa(charges_pKadet,pHs_pKadet,_acid_base[group]);
                }
            } else {
                if (this_crg<=-0.5 && last_crg>-0.5) {
                    //
                    // Get ph,charge sets and calc pKa from those
                    //
                    vector<double> pHs_pKadet;
                    vector<float> charges_pKadet;
                    int count2=count-static_cast<int>(datapoints);
                    if (count2<0) {count2=0;}
                    for (double pH2=max(pH_start,pH-datapoints*pH_step);pH2<min(pH_end,pH+datapoints*pH_step);pH2=pH2+pH_step) {
                        pHs_pKadet.push_back(pH2);
                        charges_pKadet.push_back(charges[count2][group]);
                        count2=count2+1;
                    }
                    pKa=calc_pKa(charges_pKadet,pHs_pKadet,_acid_base[group]);
                }
            }
            last_crg=this_crg;
            count=count+1;
        }
        pKas.push_back(pKa);
    }   
    //
    // We also return all the charges
    //
    // The first numbers we return is the start pH, the pH step, and the number of values
    //
    int num_pHs=0;
    for (float pH=pH_start;pH<pH_end;pH=pH+pH_step) {
        printf ("\b\b\b\b\b\b\b\b\b\b\b\bpH: %5.2f",pH);
        num_pHs++;
    }
    pKas.push_back(pH_start);
    pKas.push_back(pH_step);
    pKas.push_back(static_cast<float>(num_pHs));
    // 
    // Now add the charges
    //
    float this_crg;
    int count=0;
    for (int group=0;group<_groups;group++) {
        count=0;
        for (float pH=pH_start;pH<pH_end;pH=pH+pH_step) {
            pKas.push_back(pH);
            this_crg=charges[count][group];
            pKas.push_back(this_crg);
            count=count+1;
        }
        pKas.push_back(999.0);
        pKas.push_back(-999.0);
    }
    return pKas;
}

//
// ---------------------
//

double MC::calc_pKa(vector<float> charges,vector<double> pHs,double acid_base) {
    //
    // Calculate the pKa value from a selection of charges and pH values
    //
    // Assume perfect Henderson-Hasselbalch behaviour
    //
    // acid_base = -1.0 for acids
    // acid_base = 1.0 for bases
    //
    double ratio=0.0;
    vector<double> pKas;
    double pKa=0.0;
    int points=static_cast<int>(charges.size());
    for (int count=0;count<points;count++) {
        if (acid_base!=1.0) {
            ratio=fabs(charges[count])/(1.0-fabs(charges[count]));
        } else {
            ratio=(1.0-fabs(charges[count]))/fabs(charges[count]);
        }
        pKas.push_back(pHs[count]-log10(ratio));
    }
    //
    // Find the average of the pKa values
    //
    double sum=0.0;
    for (int count=0;count<static_cast<int>(pKas.size());count++) {
        sum=sum+pKas[count];
    }
    pKa=sum/static_cast<double>(pKas.size());
    return pKa;
}


//
// ---------------------
//

// A brief explanation of how the MC code works:

//For each pH step
//    Randomly select a residue state and evaluate the total energy of that state.
//    For each step in the Monte Carlo
//        Randomly change one residue to a different state and evaluate the total energy of that state.
//        If the new total energy is lower than the previous energy
//            Keep the new state and discard the old.
//        Otherwise if the difference (d for our purposes) is between the new and old state is less than 20.0. (No clue what the units are here.)
//            Choose a random real number [0.0, 1.0) (x for our purposes, the code calls it “tilf”)
//            If x < e^(-d)
//                Keep the new state and discard the old.
//        Otherwise keep the old state.
//
//        After 10% of the steps have been processed 
//            We assume that we have equilibrated
//            We save each resulting state after each Monte Carlo step
//    For each residue from the saved states find the probability that the residue was charged at that pH.


vector<float> MC::calc_charge(float pH) {
    //
    // Calculate the fractional charges at this pH
    //
    // Initialise random number generator
    //
    srand(time(NULL));
    //
    // Get a random starting state
    //
    vector<int> current_state;
    vector<int> try_state;
    vector<vector<int> > sum_state;
    for (int group=0;group<_groups;group++) {
        current_state.push_back(static_cast<int>(rand()%2));
        if (current_state[group]==2) {
            current_state[group]=1;
        }
        // 
        // Dummy initialisation of try_state and sum_state
        //
        try_state.push_back(0);
        //sum_state.push_back(0);
        vector<int> dummy;
        sum_state.push_back(dummy);
    }
    //
    // Get the energy of the starting state 
    //
    double current_energy=get_energy(pH,current_state);
    //
    // Need to incorporate direct switching between neutral states.. no tie now
    //
    // Start the MC loop
    //
    int eqsteps=_MCsteps/10;
    int keep=0;
    double tilf=0.0;
    double try_energy_new=0.0;
    double diff=0.0;
    for (int step=0;step<_MCsteps;step++) {
        //
        // Copy the current state to trystate
        //
        for (int count=0;count<_groups;count++) {
            try_state[count]=current_state[count];
        }
        //
        // Change a random group
        //
        int rand_group=static_cast<int>(rand()%_groups);
        //
        // Change to a random state
        //
        int rand_group_state=static_cast<int>(rand()%_num_states[rand_group]);
        // 
        // Stay in while loop until we get a different state
        //
        while (rand_group_state==current_state[rand_group]) {
            rand_group_state=static_cast<int>(rand()%_num_states[rand_group]);
        }
        // Change state
        try_state[rand_group]=abs(rand_group_state);
        //
        // Get the energy of the new state
        //
        try_energy_new=get_energy(pH,try_state);
        //
        // Keep or reject?
        //
        diff=try_energy_new-current_energy;
        if (diff<0.0) {
            //
            // Keep
            //
            for (int count=0;count<_groups;count++) {
                current_state[count]=try_state[count];
            }
            current_energy=try_energy_new;
        } else {
            if (diff<20.0) {
                tilf=static_cast<double>(rand()) / (static_cast<double>(RAND_MAX)+static_cast<double>(1));
                if (tilf<exp(-diff)) {
                    //
                    // Keep
                    //
                    for (int count=0;count<_groups;count++) {
                        current_state[count]=try_state[count];
                    }
                    current_energy=try_energy_new;
                }
            }
        }
        //
        // Record the state if we have equilibrated
        //
        if (step>eqsteps) {
            for (int count=0;count<_groups;count++) {
                sum_state[count].push_back(current_state[count]);
            }
        }
    }
    //
    // Calculate fractional charge
    //
    int sample_steps=_MCsteps-eqsteps;
    vector<float> charges_thispH;
    for (int count=0;count<_groups;count++){
        float charge=0.0;
        charge=static_cast<float>(_acid_base[count]);
        //
        // Loop over all the steps and add the correct charge
        //
        float charge_sum=0.0;
        for (int step=0;step<sum_state[count].size();step++){
            //
            // Get the charge from the charged state
            //
            int state=sum_state[count][step];
            charge_sum=charge_sum+charge*static_cast<float>(_charged_state[count][state]);
        }
        charge=charge_sum/(static_cast<float>(sample_steps));
        charges_thispH.push_back(charge);
    }
    printf ("\b\b\b\b\b\b\b\b\b\b\b\bpH: %5.2f",pH);
    fflush(stdout);
    return charges_thispH;
}

//
// --------------------
// 

double MC::get_energy(float pH,vector<int> state) {
    //
    // Calculate the energy of the present state
    //
    //printf ("\nCalculating energy\n");
    double pH_value=static_cast<double>(pH);
    double energy=0.0;
    for (int group1=0;group1<_groups;group1++) {
        //
        // Add the energy from the intrinsic pKa
        //
        int charge_grp1=_charged_state[group1][state[group1]];
        if (charge_grp1!=0) {
            energy=energy+_acid_base[group1]*lnten*(pH_value-_intpKas[group1][state[group1]]);
        } else {
            energy=energy+_intpKas[group1][state[group1]]*lnten;
        }
        //
        // Add the interaction energies
        // All energies have been corrected in pkanew, so we add all intenes
        //
        for (int group2=0;group2<_groups;group2++) {
            if (group1!=group2) {
                energy=energy+_matrix[group1][group2][state[group1]][state[group2]]/2.0;
            }
        }
    }
    return energy; 
}


