//
//  ScistErrRateInf.cpp
//
//
//  Created by Yufeng Wu on 6/20/18.
//
//

#include "ScistErrRateInf.hpp"
#include "ScistPerfPhyImp.hpp"
#include "UtilsNumerical.h"

// *************************************************************************************
// Inf error rate

const double DEF_RATE_FN_MIN = 0.000001;
const double DEF_RATE_FN_MAX = 0.5;
const double DEF_RATE_FP_MIN = 0.0000001;
const double DEF_RATE_FP_MAX = 0.05;

ScistErrRateInf ::ScistErrRateInf(ScistGenGenotypeMat &genos) : genosInput(genos), rateFNMin(DEF_RATE_FN_MIN), rateFNMax(DEF_RATE_FN_MAX), rateFPMin(DEF_RATE_FP_MIN), rateFPMax(DEF_RATE_FP_MAX), fVerbose(false)
{
    //
    rateFNOpt = 0.5 * (rateFNMin + rateFNMax);
    rateFPOpt = 0.5 * (rateFPMin + rateFPMax);
}

void ScistErrRateInf ::Infer()
{
    // EM algorithm.
    const double THRES_LARGER_RATIO = 1.05;
    double likeliMaxAll = -1.0 * HAP_MAX_INT;
    while (true)
    {
        // now search for rateFP then we are done
        std::set<std::pair<std::pair<int, int>, int>> listChangedPlaces;
        double likeliMax2 = CalcMaxProbFor(this->rateFNOpt, this->rateFPOpt, listChangedPlaces);

        if (fVerbose)
        {
            cout << "Current likelihood for optimizing false positive rate is " << likeliMax2 << ", FN estimate: " << this->rateFNOpt << ", FP estimate: " << this->rateFPOpt << endl;
        }
        if (NumericalAlgoUtils::IsLikeliSignificantlyLargeThresNum(likeliMax2, likeliMaxAll, 1, THRES_LARGER_RATIO) == false)
        {
            break;
        }
        likeliMaxAll = likeliMax2;
        UpdateEstimates(listChangedPlaces);
    }

    cout << "Optimal false negative rate is " << this->rateFNOpt << ", and optimal false positive rate is " << this->rateFPOpt << endl;
}

double ScistErrRateInf ::CalcMaxProbFor(double rateFN, double rateFP, std::set<std::pair<std::pair<int, int>, int>> &listChangedPlaces)
{
    //cout << "rateFN: " << rateFN << ", rateFP: " << rateFP << endl;
    //
    ScistGenGenotypeMat *pGenosMatTest = genosInput.Copy();

    // setup prob based on the rate
    for (int s = 0; s < genosInput.GetNumSites(); ++s)
    {
        for (int c = 0; c < genosInput.GetNumHaps(); ++c)
        {
            int allele = genosInput.GetGenotypeAt(c, s);
            double prob0 = 1.0 - rateFN;
            if (allele == 1)
            {
                prob0 = rateFP;
            }
            //cout << "Setting cell " << c << ", site " << s << ", prob0: " << prob0 << endl;

            pGenosMatTest->SetGenotypeProbAt(c, s, prob0);
        }
    }
    //cout << "Genotype matrix to test: " << endl;
    //pGenosMatTest->Dump();

    double probMax = CalcMaxProbForMat(*pGenosMatTest, listChangedPlaces);
    //cout << "For rateFN: " << rateFN << ", rateFP: " << rateFP << " CalcMaxProbFor: " << probMax << endl;

    delete pGenosMatTest;

    return probMax;
}

double ScistErrRateInf ::CalcMaxProbForMat(ScistGenGenotypeMat &genosTest, std::set<std::pair<std::pair<int, int>, int>> &listChangedPlaces)
{
    //
    ScistPerfPhyMLE phInf1(genosTest);
    phInf1.SetVerbose(false);
    phInf1.SetOutput(false);
    double res = phInf1.Infer(&listChangedPlaces);
    //cout << "In CalcMaxProbForMat: prob=" << res << ", matrix: \n";
    //genosTest.Dump();
    return res;
}

void ScistErrRateInf ::UpdateEstimates(const std::set<std::pair<std::pair<int, int>, int>> &listChangedPlaces)
{
    //
    int num0to1 = 0, num1to0 = 0;
    for (set<pair<pair<int, int>, int>>::const_iterator it = listChangedPlaces.begin(); it != listChangedPlaces.end(); ++it)
    {
        //
        if (it->second == 0)
        {
            ++num1to0;
        }
        else
        {
            ++num0to1;
        }
    }
    int num0Tot = this->genosInput.GetGenotypeNumOf(0);
    int num1Tot = this->genosInput.GetGenotypeNumOf(1);
    //cout << "In UpdateEsimate: num0to1: " << num0to1 << ", num1to0: " << num1to0 << ", num0Tot: " << num0Tot << ", num1Tot: " << num1Tot << endl;
    this->rateFNOpt = ((double)(num0to1 + 1)) / (num0to1 + num1Tot + 2);
    this->rateFPOpt = ((double)(num1to0 + 1)) / (num1to0 + num0Tot + 2);
}
