#include "P0LambdaCorrection.hh"
//////////////////////////////////////////////////////////
P0LambdaCorrection::
P0LambdaCorrection()
{
    Initialize();
}
//////////////////////////////////////////////////////////
P0LambdaCorrection::
P0LambdaCorrection( ElementContainerMatrix* dat1, ElementContainerMatrix* dat2 )
{
    Initialize();
    SetData( dat1, dat2 );
}
//////////////////////////////////////////////////////////
P0LambdaCorrection::
~P0LambdaCorrection()
{
    delete _st;
}
//////////////////////////////////////////////////////////
void P0LambdaCorrection::
Initialize(){
    _NumOfMulth = UtsusemiGetNumOfMulTh();
    _MessageTag = "P0LambdaCorrection::";
    _DAT1 = NULL;
    _DAT2 = NULL;
    _pol_lam.clear();
    _pol_val.clear();
    _st = new StringTools();
}
//////////////////////////////////////////////////////////
void P0LambdaCorrection::
SetData( ElementContainerMatrix* dat1, ElementContainerMatrix* dat2 ){
    _DAT1 = dat1;
    _DAT2 = dat2;
}
//////////////////////////////////////////////////////////
bool P0LambdaCorrection::
SetPolarizationTable( std::string _pathToData ){
    std::string foundPath=FindParamFilePath( _pathToData );
    if (foundPath==""){
        UtsusemiError( _MessageTag+"SetPolarizationTable >> No such file ("+_pathToData+")" );
        return false;
    }
    std::ifstream ifs( foundPath.c_str() );
    if (ifs.fail()){
        UtsusemiError( _MessageTag+"SetPolarizationTable >> Failed to open file ("+_pathToData+")" );
        return false;
    }
    std::string aline;
    _pol_lam.clear();
    _pol_val.clear();
    while(getline(ifs,aline)){
        if (aline.substr(0,1)!="#"){
            std::vector<std::string> conts = _st->SplitString( aline, "," );
            if (conts.size()>2){
                _pol_lam.push_back( _st->StringToDouble( conts[0] ) );
                _pol_val.push_back( _st->StringToDouble( conts[1] ) );
            }
        }
    }
    ifs.close();
    return true;
}
//////////////////////////////////////////////////////////
bool P0LambdaCorrection::
Execute( ElementContainerMatrix* _Ip, ElementContainerMatrix* _Im){
    if (_pol_lam.empty()){
        UtsusemiError(_MessageTag+"Execute >> Not read polarization table. Do SetPolarizationTable.");
        return false;
    }
    if ((_DAT1==NULL)||(_DAT2==NULL)){
        UtsusemiError(_MessageTag+"Execute >> Not set data. Do SetData.");
        return false;
    }
    if (_DAT1->PutPointer(0)->PutPointer(0)->CheckKey("Lamb")==0){
        UtsusemiError(_MessageTag+"Execute >> Given data is invalid.");
        return false;
    }

    // Interpolate polarization values with lambda values in given data
    std::vector<Double>* lambda = _DAT1->PutPointer(0)->PutPointer(0)->PutP("Lamb");
    std::vector<Double> Pn_v;
    for (UInt4 i=0; i<(lambda->size()); i++){
        Double lamb = lambda->at(i);
        Double Pn = 0.0;
        for (UInt4 j=0; j<(_pol_lam.size()-1); j++){
            if ( (lamb>=_pol_lam[j])&&(lamb<=_pol_lam[j+1]) ){
                Pn = _pol_val[j] + (_pol_val[j+1]-_pol_val[j])/(_pol_lam[j+1]-_pol_lam[j]) * (lamb - _pol_lam[j] );
                break;
            }
        }
        Pn_v.push_back(Pn);
    }

#ifdef MULTH
    omp_set_num_threads( _NumOfMulth );
#endif

    _Ip->InputHeader( _DAT1->PutHeader() );
    _Im->InputHeader( _DAT1->PutHeader() );

    UInt4 ecm_size = _DAT1->PutSize();
    _Ip->Allocate( ecm_size );
    _Im->Allocate( ecm_size );
    for (UInt4 i=0; i<ecm_size; i++){
        ElementContainerArray* eca_p = new ElementContainerArray( _DAT1->PutPointer(i)->PutHeader() );
        ElementContainerArray* eca_m = new ElementContainerArray( _DAT1->PutPointer(i)->PutHeader() );
        UInt4 eca_size = _DAT1->PutPointer(i)->PutSize();
        eca_p->Allocate( eca_size );
        eca_m->Allocate( eca_size );
        for (UInt4 j=0; j<eca_size; j++ ){
            ElementContainer* ec1 = _DAT1->PutPointer(i)->PutPointer(j);
            ElementContainer* ec2 = _DAT2->PutPointer(i)->PutPointer(j);

            std::vector<Double>* I1 = ec1->PutP( ec1->PutYKey() );
            std::vector<Double>* I2 = ec2->PutP( ec2->PutYKey() );
            std::vector<Double>* E1 = ec1->PutP( ec1->PutEKey() );
            std::vector<Double>* E2 = ec2->PutP( ec2->PutEKey() );
            std::vector<Double> Ip(I1->size(),0.0);
            std::vector<Double> Im(I1->size(),0.0);
            std::vector<Double> Ep(I1->size(),0.0);
            std::vector<Double> Em(I1->size(),0.0);

            ElementContainer* ec_p = new ElementContainer( *ec1 );
            ElementContainer* ec_m = new ElementContainer( *ec1 );
#pragma omp parallel for
#if (_OPENMP >= 200805)  // OpenMP 3.0 and later
            for (UInt4 k=0; k<(I1->size()); k++){
#else
            for (Int4 k=0; k<(I1->size()); k++){
#endif
                Double i1 = I1->at(k);
                Double i2 = I2->at(k);
                Double e1 = E1->at(k);
                Double e2 = E2->at(k);

                Double P = Pn_v[k];
                // Intensity and Error
                Ip[k] = (i1+i2)/2.0 + (i1-i2)/2.0/P;
                Im[k] = (i1+i2)/2.0 - (i1-i2)/2.0/P;
                Ep[k] = sqrt( pow(((P+1)/2.0/P*e1),2) + pow(((P-1)/2.0/P*e2),2) );
                Em[k] = sqrt( pow(((P-1)/2.0/P*e1),2) + pow(((P+1)/2.0/P*e2),2) );
            }
            // Construct Ip data
            std::string xkey = ec1->PutXKey();
            std::string ykey = ec1->PutYKey();
            std::string ekey = ec1->PutEKey();

            ec_p->Remove(ykey);
            ec_p->Remove(ekey);
            ec_p->Add(ykey, Ip);
            ec_p->Add(ekey, Ep);
            ec_p->SetKeys(xkey,ykey,ekey);
            eca_p->Set( j, ec_p );

            // Construct Im data
            ec_m->Remove(ykey);
            ec_m->Remove(ekey);
            ec_m->Add(ykey, Im);
            ec_m->Add(ekey, Em);
            ec_m->SetKeys(xkey,ykey,ekey);
            eca_m->Set( j, ec_m );
        }
        _Ip->Set( i, eca_p );
        _Im->Set( i, eca_m );
    }
    return true;
}
