#include "../include/load_lua_data.h"

/** This code looks awful, needs some serious refactoring **/
/** WARNING LUA uses tables index starting at 1, indexes must be converted **/

// TODO : Implement proper error throw & catch instead of simply quitting when a problem occurs



void load_species(lua_wrapper &state, string tableName, vector<species> &rVal)
{
    rVal.clear();
    //Gets an index of species in lua environnement
    state.open_global_table(tableName);
    state.check_type_error(LUA_TTABLE, "Loading species: ", tableName, " not found");
    
    size_t nSpe = state.table_len_top(); //Get n species

    for(size_t i = 1; i <= nSpe; i++) //Iterate on them
    {
        state.open_field_table(i);
        rVal.push_back(create_specie_top(state)); //Load specie on top
        state.close_table();
    }
    state.close_table();
    return;
}


species create_specie_top(lua_wrapper &state)
{

    //Main interface between lua Specie and C++ implementation.
    //Load infos from a species struct, from a lua index
    //Gets 8 fields :
    //id, traitId, ageClasses, nloci, nAlleleTot, ploidy, traits, parameters

    //  FIELD 1 : id -- OK
    string name = state.get_field<string>("id");

    //  FIELD 2 : traitsID -- OK
    vector<string> traitsId;
    state.open_field_table("traitsId");
    state.load_array<string>(traitsId);
    state.close_table();


    //  FIELD 3 : Age classes -- OK
    vector<int> ageClasses;
    state.open_field_table("ageClasses");
    state.load_array<int>(ageClasses);
    state.close_table();


    //FIELD 4 : nLocus -- OK
    int nLoci = state.get_field<int>("nLocus");


    //FIELD 5 : nAlleles per loci -- OK
    vector<int> nAlleles;
    state.open_field_table("nAllpLoci");
    state.load_array<int>(nAlleles);
    state.close_table();

    //FIELD 5.5 : gamma distrib

    //double gammaParam = state.get_field<double>("gammaParam");

    //FIELD 6 : ploidy -- OK
    vector<int> ploidy;
    state.open_field_table("ploidy");
    state.load_array<int>(ploidy);
    state.close_table();

    //FIELD 6.25 : chromoSplit -- OK
    vector<int> chromoSplit;
    state.open_field_table("chromoSplit");
    state.load_array<int>(chromoSplit);
    state.close_table();

    //FIELD 6.5 : selection strength matrix -- OK
    size_t nTraits = traitsId.size();
    vector<double> selectMatrix(nTraits*nTraits, 0);

    vector<string> keysSel; //key index
    state.open_field_table("selIntInteractions");
    tableref refFrames = state.table_ref_top();
    state.get_table_keys(keysSel, refFrames);
    state.open_field_table("selIntInteractions");
    bool isVal = false;

    for(string key: keysSel) // ONLY HANDLE ORDER 2 INTERACTIONS
    {
        int posSplit = key.find("_");
        string t1 = key.substr(0, posSplit);
        string t2 = key.substr(posSplit+1, key.length());
        int posT1 = -1;
        int posT2 = -1;
        for (int nt=0;nt < nTraits;nt++)
        {
            if (traitsId[nt] == t1)
            {
                posT1 = nt;
            }
            if (traitsId[nt] == t2)
            {
                posT2 = nt;
            }
        }
        if (posT1==posT2 or posT1==-1 or posT2==-1)
        {
            cout << "Interaction " << key << " can't be understood. Must be the name of two defined traits for the current species" << endl;
        }
        else
        {
            // Here we suppose that interactions are symetrical (T1->T2 = T2 -> T1)
            int interactVal;
            vector<int> interactValVec;
            interactVal = state.get_field<int>(key);
            //cout << "INTVAL " << interactVal << endl;
            selectMatrix[posT1*nTraits+posT2] = interactVal;
            selectMatrix[posT2*nTraits+posT1] = interactVal;
            isVal = true;
        }
    }
    state.close_table();

    if (isVal and false) // skipped
    {
        cout << "Generic warning: interaction value should be weaker than each direct trait selection strength"<< endl;
    }

    /*for (int i=0;i<selectMatrix.size();i++)
    {
        cout << selectMatrix[i] << " " << endl;
    }*/
        //Matrix filling should be done in regard to traits order as defined in the traits table


    //FIELD 7 loading traits
    //Ref the field & throw it to the function
    vector<trait> traits;

    state.open_field_table("traits");
    nTraits = state.table_len_top();
    for(size_t i = 1; i <= nTraits; i++) //Iterate  on each trait (subtable)
    {
        state.open_field_table(i);
        traits.push_back(get_trait_top(state, ploidy, 
                         "Loading : species : traits, field"));
        state.close_table();
    }
    state.close_table();


    //FIELD 8 : crossingRate -- OK
    vector<double> crossingRate;
    state.open_field_table("crossingRate");
    state.load_array<double>(crossingRate);
    state.close_table();


    //FIELD 9 : mutationRate -- OK
    vector<double> mutationRate;
    state.open_field_table("mutationRate");
    state.load_array<double>(mutationRate);
    state.close_table();


    //// FIELD 10 : SEX -- OK
    vector<vector<int>> inherit;
    vector<int> subInherit;
    vector<string> sexOrder = {"female", "male"};

    state.open_field_table("inheritance");

    for (auto s: sexOrder) // female then male
    {
        state.open_field_table(s);
        state.load_array<int>(subInherit);
        state.close_table();
        if (subInherit.size() < nLoci)
        {
            if (ploidy[0] == 2) // CARE, DOES NOT WORK IF PLOIDY LOOKS LIKE {2,3,3...} since check only first value
            {
                while (subInherit.size() < nLoci)
                {
                    subInherit.push_back(1);
                }
            }
            else
            {
                cout << "ERROR: sex field in Sx.lua can not be automatically completed for " << nLoci << "loci with given ploidy." << endl;
            }
        }
        inherit.push_back(subInherit);
        subInherit.clear();
    }


    state.close_table();


    //// FIELD 11 : Flux Matrix : TODO : pack this in a function
    //Ok, here we miss a time information for each frame (just an int as for now)
    vector<string> propagOrder = {"seed", "pollen"};
    vector<list<fluxMat>> fm_all;
    list<fluxMat> frames;
    vector<vector<double>> dataMat;
    /**/// size_t s = 0;

    state.open_field_table("fluxMatrix");


    // TODO : clean OR refactor this complicated part
    //Call female table then male one
    vector<string> keys; //key index


    for (auto s: propagOrder) // female then male
    {
        state.open_field_table(s);
        tableref refFrames = state.table_ref_top();
        state.open_field_table(s);
        state.get_table_keys(keys, refFrames);
        for(string key: keys)
        {
            state.open_field_table(key);
            state.load_2d_array<double>(dataMat);
            state.close_table();

            frames.push_back(fluxMat(dataMat, atoi(key.c_str())));
        }


        state.close_table(); //close female
        fm_all.push_back(frames); //Add female to a vector
        frames.clear(); //Clear temp list
    }

    state.close_table(); //Closes fluxmatrix
    

    //Cap
    int cap = state.get_field<int>("cap");
    //Growth
    double growth = state.get_field<double>("growth");
    //selfing
    double selfing = state.get_field<double>("selfingRate");    
    //maxAge
    //int maxAge = state.get_field<int>("maxAge");
    int maxAge = 1;
    //assortative mating
    string assortMating = state.get_field<string>("assortMating");
    //rho
    double rho;
    if (assortMating.length() >= 2)
    {
        cout << "Assortative mating based on trait " << assortMating;
        rho = state.get_field<double>("rho");
        cout << ", rho=" << rho << endl;
        //TODO: use a struct to pas the arguments, this is horrible
        species rv(name, traitsId, ageClasses, nLoci, nAlleles, ploidy, chromoSplit, traits,
               inherit, crossingRate, mutationRate, fm_all, selectMatrix, cap,
               growth, selfing, maxAge, assortMating, rho);

        return rv;
    }
    else
    {
        cout << "Detect no assortative mating. Construct species without AM nor rho." << endl;

        //TODO: use a struct to pas the arguments, this is horrible
        species rv(name, traitsId, ageClasses, nLoci, nAlleles, ploidy, chromoSplit, traits,
                   inherit, crossingRate, mutationRate, fm_all, selectMatrix, cap,
                   growth, selfing, maxAge, assortMating);

        return rv;
    }


}





trait get_trait_top(lua_wrapper &state, vector<int> ploidy, std::string error_id)
{
    /*
     * Main interface between lua trait and CPP representation
     * Loads the trait on top of lua stack
     * 6 fields
     * A bit of re hammering is welcome, current design look very bug prone
     * About ploidy: we need this information from species to create masks an other indexes
     */
    state.check_type_error(LUA_TTABLE, error_id, "get_trait_top: no table on top");


    // 1 : Id -- OK
    string name = state.get_field<string>("id");


    // 2 : params -- TODO FIX !
    state.open_field_table("config");
    unordered_map<string, double> parameters = state.load_unordered_table_top<double>();
    state.close_table();


    // TODO : fix "_l": is for linear => load 2d matrix instead
    // 3 : E
    vector<double> E;
    state.open_field_table("e_l");
    state.load_array<double>(E);
    state.close_table();

    //vector<vector<double>> temporalE; // SHOULD USE IMPORT LIKE state.load_unordered_table_top
    unordered_map<int, vector<double>> temporalE;

    vector<int> lEk;
    state.open_field_table("e_l_k");
    state.load_array<int>(lEk);
    state.close_table();

    vector<double> temporalEfull;
    state.open_field_table("e_l_t");
    state.load_array<double>(temporalEfull);
    state.close_table();

    size_t nEkeys = lEk.size();
    size_t subStepE = temporalEfull.size()/nEkeys;

    for (int i=0; i < nEkeys; i++)
    {
        temporalE[lEk[i]] = {};
        for (int j=0;j < subStepE;j++)
        {
            temporalE[lEk[i]].push_back(temporalEfull[i*subStepE+j]);
        }
        /*cout << lEk[i] << " : ";
        for (int l=0;l<temporalE[lEk[i]].size();l++)
        {
            cout << temporalE[lEk[i]][l] << " ";
        }
        cout << endl;*/
    }


    // 4 : Zopt
    vector<double> Zopt;
    state.open_field_table("zopt_l");
    state.load_array<double>(Zopt);
    state.close_table();

    unordered_map<int, vector<double>> temporalZopt;

    vector<int> lZoptk;
    state.open_field_table("zopt_l_k");
    state.load_array<int>(lZoptk);
    state.close_table();

    vector<double> temporalZoptfull;
    state.open_field_table("zopt_l_t");
    state.load_array<double>(temporalZoptfull);
    state.close_table();

    size_t nZoptkeys = lZoptk.size();
    size_t subStepZopt = temporalZoptfull.size()/nZoptkeys;

    for (int i=0; i < nZoptkeys; i++)
    {
        temporalZopt[lZoptk[i]] = {};
        for (int j=0;j < subStepZopt;j++)
        {
            temporalZopt[lZoptk[i]].push_back(temporalZoptfull[i*subStepZopt+j]);
        }
        /*cout << lZoptk[i] << " : ";
        for (int l=0;l<temporalZopt[lZoptk[i]].size();l++)
        {
            cout << temporalZopt[lZoptk[i]][l] << " ";
        }
        cout << endl;*/
    }


    // 5 : selInt
    vector<double> selInt;
    state.open_field_table("selInt_l");
    state.load_array<double>(selInt);
    state.close_table();

    unordered_map<int, vector<double>> temporalSelInt;

    vector<int> lSelIntk;
    state.open_field_table("selInt_l_k");
    state.load_array<int>(lSelIntk);
    state.close_table();

    vector<double> temporalSelIntfull;
    state.open_field_table("selInt_l_t");
    state.load_array<double>(temporalSelIntfull);
    state.close_table();

    size_t nSelIntkeys = lSelIntk.size();
    size_t subStepSelInt = temporalSelIntfull.size()/nSelIntkeys;

    for (int i=0; i < nSelIntkeys; i++)
    {
        temporalSelInt[lSelIntk[i]] = {};
        for (int j=0;j < subStepSelInt;j++)
        {
            temporalSelInt[lSelIntk[i]].push_back(temporalSelIntfull[i*subStepSelInt+j]);
        }
        /*cout << lSelIntk[i] << " : ";
        for (int l=0;l<temporalSelInt[lSelIntk[i]].size();l++)
        {
            cout << temporalSelInt[lSelIntk[i]][l] << " ";
        }
        cout << endl;*/
    }

    // 6 : types
    vector<ltypes> ltypesIndex;
    state.open_field_table("ltypes");
    size_t s = state.table_len_top(); //Get number of entries in the table
    for(size_t i = 1; i <= s; i++)
    {
        state.open_field_table(i);
        if (is_real_type(state))
        {
            //cout << "create the corresponding ltype." << endl;
            ltypesIndex.push_back(get_type_on_top(state, ploidy, "Type"));
        }
        /*else
        {
            cout << "top ltype has an empty allele list. Won't create the corresponding ltype." << endl;
        }*/
        state.close_table();
    }
    state.close_table();
    
    
    // nhpp
    bool nhpp = state.get_field<bool>("nhpp");
    
    
    // indexing
    bool indexed = state.get_field<bool>("indexing");


    //trait value(name, parameters, E, Zopt, ltypesIndex, selInt, nhpp, indexed);
    //NEW CONSTRUCTOR WITH temporalE
    trait value(name, parameters, temporalE, temporalZopt, ltypesIndex, temporalSelInt, nhpp, indexed);

    return value;
}


bool is_real_type(lua_wrapper &state)
{
    vector< vector <double> > all_effects;

    state.open_field_table("allelicEff");
    state.load_2d_array<double>(all_effects);
    state.close_table();
    return (all_effects.size() > 0);
}

ltypes get_type_on_top(lua_wrapper &state, vector<int> ploidy, std::string error_id)
{
    state.check_type_error(LUA_TTABLE, error_id, "get_type_on_top : no table on top");

    // 1 : id -- OK
    string name = state.get_field<string>("id");


    // 2 : allelic effects; allelicEff = {}, -- OK
    vector< vector <double> > all_effects;

    state.open_field_table("allelicEff"); 
    state.load_2d_array<double>(all_effects);
    state.close_table();
    if (all_effects.size() > 0)
    {
        // 3 : List of loci use by the type
        // WARNING : loctypes ARE indexes, they MUST be converted by substracting 1
        // Also: values sould be sorted in the lua file (correct this one day?)
        // Otherwise allelic effects are not in the correct order to be matched with a selected subset of genome
        // This is du to the fact that the binary mask does not keep the order of loci (only the value)
        vector<unsigned int> lociList;
        state.open_field_table("lociList");
        state.load_array<unsigned int>(lociList);
        for(unsigned int &i : lociList) i--; //Index lua to CPP index
        state.close_table();

        // 4 : Get power of E

        int pow = state.get_field<int>("pow");

        //construct & return the object
        ltypes value(name, all_effects, lociList, ploidy, pow);

        return value;
    }
    else
    {
        cout << "Alleff list is empty (due to input trait file), can't create the ltype object" << endl;
        exit(0);
    }
}




individual * load_ind_top(lua_wrapper &state)
{
    /*
        This function loads an individual that is on the top
        and return an initialized individual object
    */
    state.check_type_error(LUA_TTABLE, "load_ind_top failure: top is not a table");

    //STEP 1 : loading genetic information
    //Calling the relevant field
    vector<int> v_genome;
    state.open_field_table("genome");
    state.load_array<int>(v_genome);
    for(int &i : v_genome) i--; //-1 Index correction (allele numeroting from 1 in lua
    state.close_table();


    //STEP 2 : loading age
    //Calling the relevant field
    int age = state.get_field<int>("age");

    individual * value = new individual(age, v_genome);
    return value;

}




//Split in Load ind, load pop, load map?
int load_populations(lua_wrapper &state, string table_name, vector<species> & species_index, landscape & rMap)
{
/*
 * Loading intitial populations generated in lua
 * (in Lua, the table is stored as this : Species[]->Pop[]/Cells[]->Indiv
 *
 */
    population * myPop;             //Temp Population storage
    int general_count = 0;
    size_t s1, s2;

    //Gets the lua global table holding the information
    state.open_global_table(table_name); // #1
    // 1 - Iterate on species
    for (species &spe : species_index)
    {
        //We get the name of the current species to load the table we want
        state.open_field_table(spe.get_name()); // #2
        // 2 - Iterate on populations (= cells) of type it (= species)
        s1 = state.table_len_top(); //Get number of entries in the table
        for(size_t i = 1; i <= s1; i++) //Iterate on each entry
        {
            state.open_field_table(i); // #3
            myPop = new population(&spe, i-1, s1); //Allocates a new pop object
            // 3 - Iterate on each individual to fill the pop
            s2 = state.table_len_top();

            //cout << "s2" << s2 << endl;
            for(size_t j = 1; j <= s2; j++)
            {
                state.open_field_table(j); // #4
                //Loads and pushes individual on top
                myPop->add_individual(load_ind_top(state));
                general_count++;
                state.close_table(); // #4
            }
            rMap.get_cell(i-1)->add_population(myPop);
            myPop = NULL;
            state.close_table(); // #3
        }

        state.close_table(); // #2
    }

    state.close_table(); // #1

    return general_count;
}
