#include "Jikan.h"
#include "JikanDepths.h"
#include "ToJSON.h"

#include <iostream>
#include <map>
#include <string>

#ifdef INCLUDE_OPEN_MP
#include <omp.h>
#endif

using namespace std;

Jikan::Jikan()
{
    #pragma omp master
    {
        this->EventType["Synchronous events"] = set<string>();
        this->EventType["Asynchronous events"] = set<string>();

        this->JSONname = OUTPUT_NAME + string(".json");
        this->JSONdata = "";
        this->Error = "";
    }
};

#ifdef INCLUDE_MPI
    bool Jikan::ifWriteProc(MPI_Comm comm, int id)
    {
        bool if_write = false;
        int init_flag, fin_flag, iproc;

        MPI_Initialized(&init_flag);
        MPI_Finalized(&fin_flag);

        bool if_MPI = (init_flag) && (!fin_flag);
        if(if_MPI)
        {
            MPI_Comm_rank(comm, &iproc);
            if( iproc == id)
                if_write = true;
        }
        else
            if_write = true;


        return if_write;
    }
#endif

bool Jikan::ifWriteProc()
{
    return true;
}

Jikan::~Jikan()
{
    #pragma omp master
    {
        this->Events.clear();
        
        map<string, set<string> >::iterator it;

        for(it = this->EventType.begin(); it!=this->EventType.end(); ++it)
            it->second.clear();

        this->EventType.clear();
    }
}

bool Jikan::ifContains(const string& Name)
{
    map<string, EventData>::iterator it;

    it = this->Events.find(Name);
    if (it == this->Events.end())
        return false;

    return true;
}

void Jikan::Jikan_start(const string& name, const int& mode)
{
    bool ExistFlag, ifMPI, ifOpenMP;
    ExistFlag = this->ifContains(name);

    #pragma omp master
    {
        (this->Events)[name].ifStart = true;
    }

    #pragma omp barrier

    if(!ExistFlag)
    {
        #pragma omp master
        {
            (this->Events)[name] = EventData(name);
            (this->Events)[name].GetModeVals(mode);
        }
    }
    else
    {
        bool ifStart = (this->Events)[name].ifStart;

        if(!ifStart)
            return;
    }

    #pragma omp barrier
    ifMPI = (this->Events)[name].mode[0], ifOpenMP = (this->Events)[name].mode[1];

    #pragma omp master
    {
        if(!ExistFlag)
        {
            string EventTypeName;

            if(mode == TimerMode::NO_SYNC)
                EventTypeName = "Asynchronous events";
            else
                EventTypeName = "Synchronous events";

            this->EventType[EventTypeName].insert(name);
        }

        //--------------------------- MPI synchronization ---------------------------------------
        if(ifMPI)
            this->BarrierMPI();
        //---------------------------------------------------------------------------------------
    }

    //--------------------------- OpenMP synchronization --------------------------------
    if(ifOpenMP)
    {
        #pragma omp barrier 
    }
    //-----------------------------------------------------------------------------------

    #pragma omp master
    {
        this->StartEventTimer(name);
    }
}

void Jikan::Jikan_end(const string& name)
{    
    bool ExistFlag, ifStart;
    bool ifMPI, ifOpenMP, ifCUDA;

    ExistFlag = this->ifContains(name);
    ifStart = (this->Events)[name].ifStart;

    if((!ExistFlag) || (!ifStart))
        return;

    ifMPI = (this->Events)[name].mode[0], ifOpenMP = (this->Events)[name].mode[1], ifCUDA = (this->Events)[name].mode[2];

    #pragma omp master
    {
        if((ifCUDA) && (!ifMPI) && (!ifOpenMP))
        {    
            #ifdef INCLUDE_GPU_TIMER
                this->cuda_Jikan_end(name);
            #endif
        }
    }

    //--------------------------- OpenMP synchronization --------------------------------
    if(ifOpenMP)
    {
        #pragma omp barrier 
    }
    //-----------------------------------------------------------------------------------

    //--------------------------- MPI synchronization ---------------------------------------
    #pragma omp master
    {
        if(ifMPI)
            this->BarrierMPI();

        (this->Events)[name].ifStart = false;

        this->EndEventTimer(name);
    }
    //---------------------------------------------------------------------------------------
}

void Jikan::GenerateTypedOutputData(const string& EventType, string& EventTypeString)
{
    const int n = this->EventType[EventType].size();
    EventTypeString = JSON::StartTypedBlock(EventType);

    if(n == 0)
    {
        EventTypeString += "\n\t}";
        return;
    }

    string Row;
    string Name;
    set<string>::iterator it;
    set<string>::iterator end_m = this->EventType[EventType].end();
    --end_m;

    for (it = this->EventType[EventType].begin(); it != end_m; ++it)
    {
        Name = *it;
        JSON::GenerateRow(this->Events[Name], Row, JSON::SeparateRows);
        EventTypeString += Row;
    }

    Name = *end_m;
    JSON::GenerateRow(this->Events[Name], Row, JSON::FinishRows);
    EventTypeString += Row;
}

void Jikan::GenerateOutputData()
{
    map<string, set<string> >::iterator it;
    map<string, set<string> >::iterator end_m = this->EventType.end();
    --end_m;

    string Name;
    string Out = "{ ";
    string TypedOutput;

    for(it = this->EventType.begin(); it!=end_m; ++it)
    {    
        Name = it->first;
        TypedOutput = "";
        this->GenerateTypedOutputData(Name, TypedOutput);
        Out += TypedOutput + ",\n";
    }

    Name = end_m->first;
    TypedOutput = "";
    this->GenerateTypedOutputData(Name, TypedOutput);
    Out += TypedOutput + "\n}";

    this->JSONdata = Out;

    bool OutputRes = JSON::WriteOutput(this->JSONname, this->JSONdata);

    if(OutputRes != true)
    {
        this->Error = JSON::Error;
    }
}

void Jikan::SetDumpFilename(const string& name)
{
    this->JSONname = name;
}

void Jikan::StartEventTimer(const string& name)
{
    bool ifMPI = (this->Events)[name].mode[0], ifOpenMP = (this->Events)[name].mode[1], ifCUDA = (this->Events)[name].mode[2];

    if((ifCUDA))
    {
        #ifdef INCLUDE_GPU_TIMER
            this->cuda_Jikan_start(name);
        #endif
    }

    #ifdef INCLUDE_OPEN_MP
        else if( (!ifMPI) && (ifOpenMP) )
            (this->Events)[name].double_start = omp_get_wtime();
    #endif

    else
        (this->Events)[name].start = chrono::steady_clock::now();
}

void Jikan::EndEventTimer(const string& name)
{
    bool ifMPI = (this->Events)[name].mode[0], ifOpenMP = (this->Events)[name].mode[1], ifCUDA = (this->Events)[name].mode[2];

    if((ifCUDA))
    {
        #ifdef INCLUDE_GPU_TIMER
            this->cuda_Jikan_start(name);
        #endif
    }

    #ifdef INCLUDE_OPEN_MP
        else if( (!ifMPI) && (ifOpenMP))
        {
            double end = omp_get_wtime();
            double main_elapsed = end - (this->Events)[name].double_start;

            (this->Events)[name].elapsed_time += main_elapsed;
            (this->Events)[name].count ++;

            #ifdef SAVE_TIME_SERIES
                (this->Events)[name].time_series.push_back(main_elapsed);
            #endif
        }
    #endif

    else
    {
        chrono_time_type end = chrono::steady_clock::now();
        chrono_time_type start = (this->Events)[name].start;

        double main_elapsed = chrono::duration_cast<chrono::nanoseconds>(end - start).count() * 1e-9;

        (this->Events)[name].elapsed_time += main_elapsed;
        (this->Events)[name].count ++;

        #ifdef SAVE_TIME_SERIES
            (this->Events)[name].time_series.push_back(main_elapsed);
        #endif
    }
}

void Jikan::BarrierMPI()
{
#ifdef INCLUDE_MPI
    int init_flag, fin_flag;

    MPI_Initialized(&init_flag);
    MPI_Finalized(&fin_flag);

    if((!fin_flag) && init_flag)
        MPI_Barrier(MPI_COMM_WORLD);
#endif
}

class Jikan Timer;