package org.fda.dataaligner;



import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.List;
import org.fda.basicstats.BasicStatsCompute;
import org.fda.data.ReferenceFile;
import org.fda.data.ReferenceSet;
import org.fda.data.Utilities;
import org.fda.inputdataparser.FastaWriter;
import org.fda.inputdataparser.InputReader;
import org.fda.inputdataparser.ReadRecord;
import org.fda.multithreading.ScriptHelper;


/**
 *
 * @author Gokhan.Yavas
 */
public class DataPartitioner {
    private final File inputFile;
    private int numberofsubfolders;
    private final File targetFolder, alignerDir;    
    private final ReferenceSet chrset;
    private final BasicStatsCompute bsc;
    private final List<File> touchFiles;
    private ScriptHelper shelper;
    public List<File> getTouchFiles() {
        return touchFiles;
    }

    
    public DataPartitioner(File inputFile, int numberofsubfolders, File targetFolder, File alignerDir, ReferenceSet cset, BasicStatsCompute bsc){
        this.inputFile = inputFile;
        this.numberofsubfolders = numberofsubfolders;
        this.targetFolder = targetFolder;
        this.chrset = cset;
        this.alignerDir = alignerDir;
        this.bsc=bsc;
        touchFiles = new ArrayList();
    }
    public boolean divideDataAndCreateScripts(){
        int reffilenum = chrset.getReferenceFiles().size();
        File[] tmpFolders = new File[numberofsubfolders];
        File[] tmpScripts = new File[numberofsubfolders*reffilenum];
        BufferedWriter[] bwScriptArr = new BufferedWriter[numberofsubfolders*reffilenum];
        FastaWriter[] fw = new FastaWriter[numberofsubfolders];
        InputReader ir = new InputReader();
        ReadRecord f_rec;
        String s,t;
        int it_index;
        File[] tmpScriptArr;
        try{
            for(int i=0; i<numberofsubfolders; i++){
                s = Utilities.giveName(i);
                tmpFolders[i] = new File(targetFolder.getAbsolutePath()+File.separator+"datapart"+s);
                tmpFolders[i].mkdirs();                
                if(numberofsubfolders>1){                                
                    fw[i] = new FastaWriter(tmpFolders[i].getAbsolutePath()+File.separator+"datapart"+s+".fa");
                }
                for(int j=0; j<reffilenum; j++){
                    it_index = i*reffilenum + j;
                    t=Utilities.giveName(j);
                    String filename;
                    if(reffilenum!=1)
                        filename = "datapart"+s+"_to_refpart"+t;
                    else
                        filename = "datapart"+s+"_to_"+Utilities.getPrefix(chrset.getReferenceFiles().get(0).getFile());
                    tmpScripts[it_index] = new File(tmpFolders[i].getAbsolutePath()+File.separator+"s_"+filename+".sh");
                    bwScriptArr[it_index] = new BufferedWriter(new FileWriter(tmpScripts[it_index]));
                }
            }
            
            int index=0;
            int direction=1;                        
            ir.openFile(inputFile);
            
            while((f_rec=ir.readNextRecord())!=null){
                bsc.addLength(f_rec);
                // time to write this read into a file
                if(numberofsubfolders>1){
                    if(index<numberofsubfolders && index >= 0){                    
                        fw[index].write(f_rec);
                        index = index + direction;
                    }
                    else{
                        direction = direction * -1;
                        index = index + direction;                    
                        fw[index].write(f_rec);
                        index = index + direction;
                    }                 
                }
            }
            ir.closeReader();
            if(numberofsubfolders>1){
                for(int g = 0; g<fw.length; g++){                
                    fw[g].closeWriter();
                }
            }
            long actualNumberOfSubfolders = ir.getReadRecordNumber();
            if(actualNumberOfSubfolders<numberofsubfolders){
                // now we have to delete the excess folder that were not supposed to be created in the first place                
                for(int i=(int)actualNumberOfSubfolders; i< numberofsubfolders;i++){
                    for(int j=0; j<reffilenum; j++){
                        it_index = i*reffilenum + j;
                        bwScriptArr[it_index].close();
                    }
                    Utilities.deleteFolder(tmpFolders[i]);
                }
                numberofsubfolders = (int)actualNumberOfSubfolders;
                
                tmpScriptArr = new File[numberofsubfolders*reffilenum];
                for(int i=0; i<numberofsubfolders; i++){             
                    for(int j=0; j<reffilenum; j++){
                        it_index = i*reffilenum + j;
                        tmpScriptArr[it_index]=tmpScripts[it_index];
                    }
                }
                tmpScripts = tmpScriptArr;
            }
            ReferenceFile c;
            String chr ;
            for(int i=0; i<numberofsubfolders; i++){             
                for(int j=0; j<reffilenum; j++){
                    it_index = i*reffilenum + j;
                    bwScriptArr[it_index].append("#!/bin/bash"+Utilities.ls+"#$ -e "+tmpFolders[i].getAbsolutePath()+Utilities.ls+"#$ -o "+tmpFolders[i].getAbsolutePath()+
                            Utilities.ls+"#$ -S /bin/bash"+Utilities.ls);
                    bwScriptArr[it_index].append("hostname"+Utilities.ls);
                    bwScriptArr[it_index].append("cd "+tmpFolders[i].getAbsolutePath()+Utilities.ls);
                    c = this.chrset.getReferenceFiles().get(j);                    
                    chr = Utilities.getPrefix(c.getFile());
                    //c.getPrefix();      
                    
                    String outputname; 
                    if(numberofsubfolders==1){                        
                        outputname = Utilities.getPrefix(inputFile)+"_to_"+chr;
                        switch(Utilities.alignmentTool){                            
                            case nucmer:
                                //bwScriptArr[it_index].append(alignerDir.getAbsolutePath()+File.separator+"nucmer --mum -p "+chr+" "+c.getAbsolutePath()+" "+inputFile.getAbsolutePath()+Utilities.ls);
                                bwScriptArr[it_index].append(alignerDir.getAbsolutePath()+File.separator+"nucmer --mum -p "+outputname+" "+c.getAbsolutePath()+" "+inputFile.getAbsolutePath()+Utilities.ls);
                                break;
                            case minimap2:
                                bwScriptArr[it_index].append(alignerDir.getAbsolutePath()+File.separator+"minimap2 -c -x asm5 --mask-level 0.9 --min-occ 200 -g 2500 --score-N 2 "
                                        +c.getAbsolutePath()+" "+inputFile.getAbsolutePath()+" -t "+Utilities.no_threads+" > "+outputname+".paf"+Utilities.ls);
                                //        +c.getAbsolutePath()+" "+inputFile.getAbsolutePath()+" | tee "+outputname+".paf"+Utilities.ls);
                                break;
                        }
                    }
                    else{
                        outputname = Utilities.getPrefix(fw[i].getFile())+"_to_"+chr;
                        switch(Utilities.alignmentTool){
                            case nucmer:
                                //bwScriptArr[it_index].append(alignerDir.getAbsolutePath()+File.separator+"nucmer --mum -p "+chr+" "+c.getAbsolutePath()+" "+fw[i].getFile().getAbsolutePath()+Utilities.ls);
                                bwScriptArr[it_index].append(alignerDir.getAbsolutePath()+File.separator+"nucmer --mum -p "+outputname+" "+c.getAbsolutePath()+" "+fw[i].getFile().getAbsolutePath()+Utilities.ls);
                                break;
                            case minimap2:
                                bwScriptArr[it_index].append(alignerDir.getAbsolutePath()+File.separator+"minimap2 -c -x asm5 --mask-level 0.9 --min-occ 200 -g 2500 --score-N 2 "
                                        +c.getAbsolutePath()+" "+fw[i].getFile().getAbsolutePath()+" -t "+Utilities.no_threads+" > "+outputname+".paf"+Utilities.ls);
                                //        +c.getAbsolutePath()+" "+fw[i].getFile().getAbsolutePath()+" | tee "+outputname+".paf"+Utilities.ls);break;                                
                                break;
                        }
                        
                    }
                    //touchFiles.add(new File(tmpFolders[i].getAbsolutePath()+File.separator+chr+".touch"));
                    touchFiles.add(new File(tmpFolders[i].getAbsolutePath()+File.separator+outputname+".touch"));
                    bwScriptArr[it_index].append("touch "+touchFiles.get(touchFiles.size()-1).getAbsolutePath()+Utilities.ls);
                    bwScriptArr[it_index].close();                                
                }                
            }
            
            shelper = new ScriptHelper(tmpScripts);
            shelper.run();
            return true;            
        }
        catch(Exception e){e.printStackTrace(); return false;}
    }       
}
