/////////////////////////////////////////
/*
Supplementary material to the manuscript "Site-specific recombinatorics: in situ cellular barcoding with the Cre Lox system".
Authors: Tom S. Weber, Mark Dukes, Denise Miles, Stefan Glaser, Shalin Naik and Ken R. Duffy
for publication in "BMC Systems Biology"

The present code computes barcode probabilities (Fig. 3 C) and average number of inversions and excisions (Fig. 3 D) for a
Lox barcoding cassette with m Lox sites.
With  modifications (specified below) it also computes the data for Fig. 3 E and Fig. 4 E (inset) and distributions shown in Fig. 4 D-F.

compilation: g++ -O4 Additional\ File\ 1.cpp -o LoxCodes
run:         ./LoxCodes
*/
//////////////////////////////////////////

//headers
#include <iostream>
#include <iomanip>
#include <fstream>
#include <map>
#include <set>
#include <random>
using namespace std;

//global random number generator
std::random_device g;
std::mt19937_64 R(g());

struct halfLox
/*
one half of a Lox site
self :  index
left :  left neighbour index
right:  right neighbour index
color:  0 or 1 (gray or black in Fig. 1 A) defining orientation of the Lox site
*/
{
  int self,left,right,color;
  halfLox(int self, int left, int right, int color) :self(self), left(left),right(right),color(color) {};
};

//comparison operator
bool operator==(const halfLox& lhs, const halfLox& rhs)
{
    return (lhs.self   == rhs.self &&
            lhs.left   == rhs.left &&
            lhs.right   == rhs.right &&
            lhs.color   == rhs.color);
}

//ordering operator
bool operator<(const halfLox& lhs, const halfLox& rhs)
{
    return lhs.self < rhs.self;
}


//function performing recombination either leading to inversion or excision depending on the orientation of the Lox sites involved (black or gray, see Fig. 1 A)
void recom(int a, int b, std::vector<halfLox> & cass)
{
   if(cass[a].color!=cass[b].color) //inversion
   {
       if(cass[a].left==cass[a+1].self) cass[a].left=cass[b].self ; if(cass[a].right==cass[a+1].self) cass[a].right=cass[b].self ;
       if(cass[a+1].left==cass[a].self) cass[a+1].left=cass[b+1].self ; if(cass[a+1].right==cass[a].self) cass[a+1].right=cass[b+1].self ;
       
       if(cass[b].left==cass[b+1].self) cass[b].left=cass[a].self ; if(cass[b].right==cass[b+1].self) cass[b].right=cass[a].self ;
       if(cass[b+1].left==cass[b].self) cass[b+1].left=cass[a+1].self ; if(cass[b+1].right==cass[b].self) cass[b+1].right=cass[a+1].self ;
   }
    
   else //excision
   {
       if(cass[a].left==cass[a+1].self) cass[a].left=cass[b+1].self ; if(cass[a].right==cass[a+1].self) cass[a].right=cass[b+1].self ;
       if(cass[a+1].left==cass[a].self) cass[a+1].left=cass[b].self ; if(cass[a+1].right==cass[a].self) cass[a+1].right=cass[b].self ;
       
       if(cass[b].left==cass[b+1].self) cass[b].left=cass[a+1].self ; if(cass[b].right==cass[b+1].self) cass[b].right=cass[a+1].self ;
       if(cass[b+1].left==cass[b].self) cass[b+1].left=cass[a].self ; if(cass[b+1].right==cass[b].self) cass[b+1].right=cass[a].self ;
   }

};

//recursive function to traverse the elements of the cassette after recombination, reordering and removing excised parts
void traverse(int a,int end, std::vector<halfLox> & cass, std::vector<halfLox> & new_cass,std::vector<int> & new_code)
{
    if(a==0) {new_cass.clear(); new_code.clear();}
    if(cass[a].right==end) {new_cass.push_back(cass[a]);new_code.push_back(cass[a].self); return;}
    
    bool right=true,left=false;
    
    if(new_cass.size()!=0)
    {
        if(cass[a].right==new_cass.back().self){left=true; right=false;}
    }
    
    int next=0; for(int i=0;i<cass.size();i++) if((right==true && cass[a].right==cass[i].self) || (left==true && cass[a].left==cass[i].self)) next=i;
    new_cass.push_back(cass[a]);
    //new_code is an array of indexes
    new_code.push_back(cass[a].self);
    traverse(next,end, cass, new_cass, new_code);
}


//Barcodes with 3 elements still undergo inversions and can appear in two different configurations (would thus be counted twice).
//Function brings the two possible configuration into a single 'standard' configuration.
void standard(std::vector<int> & code)
{
    if(code.size()!=8) return;
    std::vector<int>inv_code(code);
    for(int i=1; i<code.size()-1;i++) inv_code[i]=code[code.size()-1-i];
    if(code[1]>inv_code[1]) code=inv_code;
}


int main()
{
    //data structure to store barcodes
    std::set<std::vector<int> > codes;
    
    //parameters defining the cassette
    std::vector<halfLox> ini_cass, cass, new_cass;std::vector<int> code;
    int m=14; //number of Lox sites
    int n=m*2; //number of half Lox sites
    int e=m-1;//number of code elements
    int color[]={0,1, 1,0, 0,1, 1,0, 0,1, 1,0, 0,1, 1,0, 0,1, 1,0, 0,1, 1,0, 0,1, 1,0}; //color of the half Lox sites (gray or black, see Fig. 1 A)
    
    //initializing the cassette with n Lox site halves, which corresponds to m Lox sites and (m-1) code elements.
    for(int i=0;i<n;i++)
    {
        int left=i-1, right=i+1;
        halfLox e(i,left,right,color[i]);
        ini_cass.push_back(e);
    }

    // brute force loop to find all possible barcodes
    for(int rounds=0; ; rounds++)
    {
        cass=ini_cass; //we start with full cassette
      //recombination until we reach a stable configuration
    for(;;)
    {
        //random number generator for choosing uniformly from the cass.size()/2 Lox sites
        std::uniform_int_distribution<>uni(0,cass.size()/2-1);
        //choosing 2 Lox sites at random
        int Lox1=uni(R),Lox2=uni(R);
        //if they are too close no interaction
        if(abs(Lox1-Lox2)<3) continue;
        //recombination
        recom(2*std::min(Lox1,Lox2), 2*std::max(Lox1,Lox2), cass);
        //cleaning up, reordering, and removing excised elements
        traverse(0,n,cass,new_cass,code);
        
        cass=new_cass;
        
        //we stop if final size is reached (4 and 2 Lox sites or 8 and 4 half Lox sites)
        if(cass.size()/2==4 || cass.size()/2==2) break;
    }
        //standardizing codes with 4 Lox sites
        standard(code);
        //add code to the set of codes. If the set already contains the barcode, by definition nothing happens.
        codes.insert(code);
        
        //we stop if we have found all codes
        if(codes.size()==pow(e-1,2)*(e+1)/2+(e+1)) break;
    }
    
    //creating a map that associates each barcode with a unique integer
    std::map<std::vector<int>,int> MAP;std::set<std::vector<int> >::iterator itc;
    int j=0; for(itc=codes.begin(); itc!=codes.end(); itc++,j++) MAP[(*itc)]=j;
    
    //data structure to count how often a specific barcode appears, and how many inversions and excisions were needed to create it
    std::vector<int>histo(codes.size(),0),steps_inv(codes.size(),0),steps_ex(codes.size(),0);
    
    ///////////computing barcode probabilities using 1e8 independent Monte Carlo simulations//////////////
    for(int rounds=0; rounds<1e8; rounds++)
    {
        //we start with full cassette
        cass=ini_cass;
        
        int count_inv=0;int count_ex=0;
        
        //recombination until reaching size-stable barcode
        for(;;)
        {
            //random number generator for choosing uniformly from the cass.size()/2 Lox sites
            std::uniform_int_distribution<>uni(0,cass.size()/2-1);
            //selecting two Lox sites uniformly
            int Lox1=uni(R),Lox2=uni(R);
            
            //if Lox sites are sufficiently apart, apply recombination
            if(abs(Lox1-Lox2)>=3)
            {
                recom(2*std::min(Lox1,Lox2), 2*std::max(Lox1,Lox2), cass);
                if(cass[2*std::min(Lox1,Lox2)].color!=cass[2*std::max(Lox1,Lox2)].color) count_inv++;
                else count_ex++;
            }
        
            //cleaning up, reordering, and removing excised elements
            traverse(0,n,cass,new_cass,code);
            
            cass=new_cass;
            
            ///we stop if final size is reached (4 and 2 Lox sites or 8 and 4 half Lox sites)
            if(cass.size()/2==4 || cass.size()/2==2) break;
            
        }
        //standardizing codes with 4 Lox sites
        standard(code);
        
        //increase the respective counts in the histograms
        histo[MAP[code]]++;steps_inv[MAP[code]]+=count_inv; steps_ex[MAP[code]]+=count_ex;
    }
    
    
    
    //writing results to file
    ofstream file1("barcode_probabilities.txt");
    file1<<"barcode.probability avg.inversions avg.excisions"<<endl;
    for(int i=0; i<histo.size(); i++) file1<<fixed<<setprecision(10)<<i<<" "<<setw(10)<<histo[i]/1e8<<" "<<setw(10)<<(double)steps_inv[i]/histo[i]<<" "<<setw(10)<<(double)steps_ex[i]/histo[i]<<endl;
    file1.close();
    return 0;
}

/*
///Fig. 4 D: 2 simultaneous Lox-Lox interactions:

//replace line 184-192 by
//////////////////////////////////////////////////////////////////
 //we choose 4 Lox sites at random, 1 and 2 interact, and 3 and 4.
 int Lox1=uni(R), Lox2=uni(R),Lox3=uni(R),Lox4=uni(R);
 //all sites must be distinct
 if(Lox1==Lox3 || Lox1==Lox4 || Lox2==Lox3 || Lox2==Lox4) continue;
 //recombination
 if(abs(Lox1-Lox2)>=3) recom(2*std::min(Lox1,Lox2), 2*std::max(Lox1,Lox2), cass);
 if(abs(Lox3-Lox4)>=3) recom(2*std::min(Lox3,Lox4), 2*std::max(Lox3,Lox4), cass);
 ///////////////////////////////////////////////////////////////////
 
///Fig. 4 E: transient Cre expression:

 //replace line 133-218 by
 //////////////////////////////////////////////////////////////////
 //data structures to store barcodes from transient Cre expression
 std::multiset<std::vector<int> > tcodes; //will store all barcodes
 std::set<std::vector<int> > tunique;     //will store a given barcode only once
 std::set<std::vector<int> >::iterator itu;
 
 //generate 1e6 random codes with Poisson distributed number of recombination events
 /////////////////////////////////////
 for(int rounds=0; rounds<1e6 ; rounds++)
 {
 cass=ini_cass;
 std::poisson_distribution<>P(1); //Poisson distribution with expected value 1
 for(int k=0;k<P(R);k++)
 {
 std::uniform_int_distribution<>uni(0,cass.size()/2-1);
 int Lox1=uni(R),Lox2=uni(R);
 if(abs(Lox1-Lox2)<3) continue;
 recom(2*std::min(Lox1,Lox2), 2*std::max(Lox1,Lox2), cass);
 traverse(0,n,cass,new_cass,code);
 cass=new_cass;
 if(cass.size()/2==4 || cass.size()/2==2) break;
 }
 tcodes.insert(code);
 tunique.insert(code);//only inserts if code is not in tunique yet
 }
 
 //count how often a code was created
 ofstream file2("transient_barcode_probabilities.txt");
 for(itu=tunique.begin(); itu!=tunique.end();itu++) file2<<tcodes.count((*itu))<<endl;
 file2.close();
//////////////////////////////////////////////////
 
///Fig. 4 F: distance dependent Lox-Lox interactions:
 
 //penalize sites that are further apart by adding in line 185
 /////////////////////////////////////////////////////////////
 std::uniform_real_distribution<>U(0,1);
 if(U(R)>1.0/abs(Lox1-Lox2)) continue;
 /////////////////////////////////////////////////////////////
 
 
///Fig. 3 E: 99%-unique barcodes for 3 cassettes

 //to compute 99-unique barcodes add the following after line 210
 //////////////////////////////////////////////////////////////
 std::vector<double>PDF;
 //product distribution
 for(int i=0; i<histo.size(); i++)for(int j=0; j<histo.size(); j++)for(int k=0; k<histo.size(); k++) {PDF.push_back(histo[i]*histo[j]*histo[k]*pow(1e8,-3));}
 std::sort(PDF.begin(),PDF.end());
 
 //99%-unique barcodes
 double f=0.99;
 double sumP=0, sumP2=0;
 ofstream file3("99_unique_barcodes.txt");
 for(int i=0; i<PDF.size();i++)
 {
 double p=PDF[i];
 sumP+=p;                            //total probability mass of selected barcodes
 sumP2+=p*p;                         //sum of squares
 sumP2_scaled=sumP2/pow(sumP,2);     //normalized
 double j=(1-f)/sumP2_scaled+1.0;    //total number of codes that are generated (j that solves Eq. 2)
 
 double f_unique_barcodes=j;
 double induced_barcodes=j/sumP;
 file3<<induced_barcodes<<" "<<f_unique_barcodes<<endl;
 }
 file3.close();
 ///////////////////////////////////////////////
*/