/** Program demonstracyjny algorytmu EM dla obl. prawdopodobienstwa haplotypow z 'niemymi' allelami.
    Implementacja jedynie w celu ilustracji.
*/

#include <iostream>
#include <numeric>
#include <iterator>
#include <algorithm>
#include <vector>

using namespace std;

const char* HELLO_MESSAGE =
"Program demonstracyjny, oblicza prawdopodobienstwo haplotypow dla nastepujacego przypadku:\n"
"2 loci, kazde z nich 2 allele, oznaczane 'A' 'a' 'B' 'b'. 'Niemy' allel to 'a'\n"
"Dla tego przypadku mozliwe sa nastepujace genotypy i rozklady haplotypow, ktore je tworza\n"
"-------------------------------------------------------------------\n"
"   GENOTYP       |    ROZKLADY         | LICZBA ROZNYCH ROZKLADOW  \n"
"-------------------------------------------------------------------\n"
"  AB             | AB/AB, AB/aB        |   2\n"
"  ABb            | AB/Ab, AB/ab, Ab/aB |   3\n"
"  Ab             | Ab/Ab, Ab/ab        |   2\n"
"  aB             | aB/aB               |   1\n"
"  aBb            | aB/ab               |   1\n"
"  ab             | ab/ab               |   1\n"
"-------------------------------------------------------------------\n";

/** w przykladzie wystepuja nastepujace genotypy (to co obserwujemy) */
const char* GENOTYPE [] = { "AB", "ABb", "Ab", "aB", "aBb", "ab" };
const int GENOTYPE_NUM = sizeof(GENOTYPE) / sizeof(GENOTYPE[0]);

/** roznych rozkladow haplotypow (dla 2 loci, kazdy z nich 2 allele) jest 10 */
const char* RESOLUTION [] = { "AB/AB", "AB/Ab", "AB/aB", "AB/ab", "Ab/Ab", "Ab/aB", "Ab/ab", "aB/aB", "aB/ab", "ab/ab" };
const int RESOLUTION_NUM = sizeof(RESOLUTION) / sizeof(RESOLUTION[0]);

/** roznych haplotypow jest 4 */
const char* HAPLOTYPE[] = { "AB", "Ab", "aB", "ab" };
const int HAPLOTYPE_NUM = sizeof(HAPLOTYPE) / sizeof( HAPLOTYPE[0]);

/** przyklad wygenoerowany sztucznie, dla p(AB) = 0.2, p(Ab) = 0.2, p(aB) = 0.2, p(ab) = 0.4
    Wtedy najwieksze prowdopodobienstwo ma obserwacja nastepujacej liczby odpowiednich genotypow */
const int GENOTYPE_FREQ_EXAMPLE[] = { 3, 8, 5, 1, 4, 4 };

/** stale przydatne do inicjacji prawdopodobienstw genotypow. */
const double ONE = 1.0;
const double SEC = 0.5;
const double THR = 1.0 / 3.0;

/** liczba krokow algorytmu */
const int STEPS = 5;

namespace {

    template<class T>
    struct ColectionDescription : public vector<T> {
        /** tablica napisow opisujaca poszczegolne elementy wektora (ma dlugosc size) */
        const char** desc_;

        ColectionDescription(int size, const char* desc[])
            : vector<T>(size), desc_(desc) {}

    };

    template<class T>
    std::ostream& operator<<(std::ostream& os, const ColectionDescription<T>& prob) {
        os.precision(2);
        for(unsigned int i=0; i< prob.size(); i++) {
            os << prob.desc_[i] << '(' << prob.at(i) << "),";
        }
        return os;
    }

    typedef ColectionDescription<int> Frequences;
    typedef ColectionDescription<double> Probability;
}

/** oblicza prawd. genotypow na podstawie: poprzednich szacowan prawdpopodobienstw rozkladow haplotypow i genotypow,
    czestosci poszczegolnych genotypow oraz liczby przykladow */
Probability calcResolution(const Probability& resolution, const Probability& genotype, const Frequences& freq, int num) {
    /** Prawdopodobienstawa sa obliczane z nastepujacych wzorow
        P(AB/AB)' = n(AB)/n * P(AB/AB)/P(AB)
        P(AB/Ab)' = n(ABb)/n * P(AB/Ab)/P(ABb)
        P(AB/aB)' = n(AB)/n * P(AB/aB)/P(AB)
        P(AB/ab)' = n(ABb)/n * P(AB/ab)/P(ABb)
        P(Ab/Ab)' = n(Ab)/n * P(Ab/Ab)/P(Ab)
        P(Ab/aB)' = n(ABb)/n * P(Ab/aB)/P(ABb)
        P(Ab/ab)' = n(Ab)/n * P(Ab/ab)/P(Ab)
        P(aB/aB)' = n(aB)/n
        P(aB/ab)' = n(aBb)/n
        P(ab/ab)' = n(ab)/n
        gdzie:
        RESOLUTION  = { "AB/AB", "AB/Ab", "AB/aB", "AB/ab", "Ab/Ab", "Ab/aB", "Ab/ab", "aB/aB", "aB/ab", "ab/ab" };
        GENOTYPE = { "AB", "ABb", "Ab", "aB", "aBb", "ab" };
    */
    const int GENO_FOR_RESOLUTION[] = { 0, 1, 0, 1, 2, 1, 2, 3, 4, 5 };

    Probability ngen(RESOLUTION_NUM, RESOLUTION);
    for(unsigned int i = 0; i < ngen.size(); i++ ) {
        int geno_idx = GENO_FOR_RESOLUTION[i];
        ngen[i] = freq[geno_idx] / (double)num * resolution[i]/genotype[geno_idx];
    }
    return ngen;
}

/** oblicza prawd. haplotypow na podst. genotypow - zlicza geny */
Probability calcHaplotypeFromResolution(const Probability& resolution) {
    /** licze wystepowanie haplotypow w odp. genotypach
        2 * P(AB) = P(AB/AB) + P(AB/AB) + P(AB/Ab) + P(AB/aB) + P(AB/ab)
        2 * P(Ab) = P(AB/Ab) + P(Ab/Ab) + P(Ab/Ab) + P(Ab/aB) + P(Ab/ab)
        2 * P(aB) = P(AB/aB) + P(Ab/aB) + P(aB/aB) + P(aB/aB) + P(aB/ab)
        2 * P(ab) = P(AB/ab) + P(Ab/ab) + P(aB/ab) + P(ab/ab) + P(ab/ab)

        suma zawiera indeksy gdy:
        RESOLUTION  = { "AB/AB", "AB/Ab", "AB/aB", "AB/ab", "Ab/Ab", "Ab/aB", "Ab/ab", "aB/aB", "aB/ab", "ab/ab" };
        HAPLOTYPE = { "AB", "Ab", "aB", "ab" };
    */

    Probability haplotypy(HAPLOTYPE_NUM, HAPLOTYPE);
    haplotypy[0] = ( 2 * resolution[0] + resolution[1] + resolution[2] + resolution[3] ) / 2.0;
    haplotypy[1] = ( 2 * resolution[4] + resolution[1] + resolution[5] + resolution[6] ) / 2.0;
    haplotypy[2] = ( 2 * resolution[7] + resolution[2] + resolution[5] + resolution[8] ) / 2.0;
    haplotypy[3] = ( 2 * resolution[9] + resolution[3] + resolution[6] + resolution[8] ) / 2.0;
    return haplotypy;
}


/** oblicza prawd. genotypow na podstawie prawd. haplotypow:
    RESOLUTION  = { "AB/AB", "AB/Ab", "AB/aB", "AB/ab", "Ab/Ab", "Ab/aB", "Ab/ab", "aB/aB", "aB/ab", "ab/ab" };
    HAPLOTYPE = { "AB", "Ab", "aB", "ab" };
*/
Probability calcResolutionFromHaplotype(const Probability& haplotype) {
    /** tabele wspolczynnikow sa inicjowane zgodnie z ponizszymi regulami obliczania prawdopodobienstwa

    P(AB/AB) = 1*P(AB)*P(AB)         P(Ab/Ab) = 1*P(Ab)*P(Ab)
    P(AB/Ab) = 2*P(AB)*P(Ab)         P(Ab/aB) = 2*P(Ab)*P(aB)
    P(AB/aB) = 2*P(AB)*P(aB)         P(Ab/ab) = 2*P(Ab)*P(ab)
    P(AB/ab) = 2*P(AB)*P(ab)

    P(aB/aB) = 1*P(aB)*P(aB)         P(ab/ab) = 1*P(ab)*P(ab)
    P(aB/ab) = 2*P(aB)*P(ab)

    */
    const int SIZE = 2;
    const int SKL[RESOLUTION_NUM][SIZE] =
        { {0,0}, {0,1}, {0,2}, {0,3}, {1,1}, {1,2}, {1,3}, {2,2}, {2,3}, {3,3} };
    const int WSP[RESOLUTION_NUM] = { 1, 2, 2, 2, 1, 2, 2, 1, 2, 1 };

    Probability resolution(RESOLUTION_NUM, RESOLUTION);
    for(unsigned int i=0; i < resolution.size(); i++ ) {
        resolution[i] = haplotype[ SKL[i][0] ]*haplotype[ SKL[i][1] ]*WSP[i];
    }
    return resolution;
}

/** oblicza prawd. genotypow na podstawie prawd. rozkladow haplotypow
 */
Probability calcGenotypeFromResolution(const Probability& resolution) {

    /** Prawdopodobienstawa sa obliczane z nastepujacych wzorow
        P(AB) = P(AB/AB)+P(AB/aB)
        P(ABb) = P(AB/Ab)+P(AB/ab)+P(Ab/aB)
        P(Ab) = P(Ab/Ab)+P(Ab/ab)
        P(aB) = P(aB/aB)
        P(aBb) = P(aB/ab)
        P(ab) = P(ab/ab)

        gdzie:
        RESOLUTION  = { "AB/AB", "AB/Ab", "AB/aB", "AB/ab", "Ab/Ab", "Ab/aB", "Ab/ab", "aB/aB", "aB/ab", "ab/ab" };
        GENOTYPE = { "AB", "ABb", "Ab", "aB", "aBb", "ab" };

    */
    Probability genotype(GENOTYPE_NUM, GENOTYPE);
    genotype[0] = resolution[0] + resolution[2];
    genotype[1] = resolution[1] + resolution[3] + resolution[5];
    genotype[2] = resolution[4] + resolution[6];
    genotype[3] = resolution[7];
    genotype[4] = resolution[8];
    genotype[5] = resolution[9];
    return genotype;
}


int main() {
    Frequences freq(GENOTYPE_NUM, GENOTYPE);
    copy( GENOTYPE_FREQ_EXAMPLE, GENOTYPE_FREQ_EXAMPLE + GENOTYPE_NUM, freq.begin() );
    int num_examples = accumulate(freq.begin(), freq.end(), 0);
    cout << HELLO_MESSAGE;
    cout << "Przykladowe obserwacje dla " << num_examples << " nastepujacych przykladow:" << endl;
    cout << freq << endl;

    /** Przy inicjacji zakladam, ze prawdopodobienstwo jest rowno rozlozone dla wszystkich rozkladow haplotypow
        tworzacych dany genotyp, wiec liczba ta jest rowna odwrotnosci liczby rozkladow dla danego genotypu */
    const double RESOLUTION_INIT[RESOLUTION_NUM] = { SEC, THR, SEC, THR, SEC, THR, SEC, ONE, ONE, ONE };
    Probability resolution(RESOLUTION_NUM, RESOLUTION);
    copy( RESOLUTION_INIT, RESOLUTION_INIT + RESOLUTION_NUM, resolution.begin() );

    /** inicjuje wartoscia neutralna rowna 1 (kazdy genotyp ma prawdopodobienstwo rowne 1 na poczatu */
    Probability genotype(GENOTYPE_NUM, GENOTYPE);
    fill( genotype.begin(), genotype.end(), 1.0 );

    /** inicjuje rownymi prawdopodobienstwami */
    Probability haplotype(HAPLOTYPE_NUM, HAPLOTYPE);
    fill( haplotype.begin(), haplotype.end(), 1.0 / HAPLOTYPE_NUM );

	for(int j = 0; j < STEPS; j++) {
		cout << "STEP " << j << endl;
        cout << "ROZKLAD: " << resolution << endl;
        cout << "GENOTYPY: " << genotype << endl;
        resolution = calcResolution(resolution, genotype, freq, num_examples );
        cout << "ROZKLAD': " << resolution << endl;
        haplotype = calcHaplotypeFromResolution(resolution);
        cout << "HAPLOTYPY: " << haplotype << endl;
        resolution = calcResolutionFromHaplotype(haplotype);
        genotype = calcGenotypeFromResolution(resolution);
	}
    return 0;
}