package org.fda.main;

import com.beust.jcommander.JCommander;
import com.beust.jcommander.ParameterException;
import java.io.BufferedInputStream;
import java.io.EOFException;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.net.URISyntaxException;
import java.nio.file.Files;
import java.nio.file.StandardCopyOption;
import java.security.CodeSource;
import java.util.Scanner;
import java.util.Set;
import java.util.TreeSet;
import org.fda.commands.AlignEvalCommand;
import org.fda.data.ReferenceSet;
import org.fda.data.Utilities;
import org.fda.dataaligner.DataAligner;
import org.fda.commands.ReevaluateCommand;
import org.fda.commands.ChrPickerCommand;
import org.fda.commands.ComputeModelCommand;
import org.fda.commands.ContigGenerateCommand1;
import org.fda.commands.ContigGenerateCommand2;
import org.fda.commands.ContigGenerateCommand3;
import org.fda.commands.ContigGenerateCommand4;
import org.fda.commands.ContigGenerateCommand5;
import org.fda.commands.FilterCommand;
import org.fda.commands.MisassembleCommand;
import org.fda.commands.RecomputeModelCommand;
import org.fda.commands.ShowModelCommand;
import org.fda.contiggenerator.ContigGenerator1;
import org.fda.contiggenerator.ContigGenerator2;
import org.fda.contiggenerator.ContigGenerator3;
import org.fda.contiggenerator.ContigGenerator4;
import org.fda.contiggenerator.ContigGenerator5;
import org.fda.contiggenerator.MisassemblyGenerator;
import org.fda.data.Enums;
import org.fda.evaluator.Evaluator;
import org.fda.evaluator.Reevaluator;
import org.fda.evaluator.Reevaluator4Artifact;
import org.fda.functionality.ShowModel;
import org.fda.inputdataparser.ChromosomePicker;
import org.fda.inputdataparser.FastaFilter;
import org.fda.regression.ModelCompute;
import org.fda.regression.RegressionModel;
/**
 *
 * @author Gokhan.Yavas
 */
public class DNAQET {
    private static final String reeval_cmd_string = "reevaluate";
    private static final String align_eval_cmd_string = "evaluate";
    private static final String generate_cmd_string1 = "generate1";
    private static final String generate_cmd_string2 = "generate2";
    private static final String generate_cmd_string3 = "generate3";
    private static final String generate_cmd_string4 = "generate4";
    private static final String generate_cmd_string5 = "generate5";
    private static final String compmodel_cmd_string = "model";
    private static final String chrpick_cmd_string = "chrpick";
    private static final String recompmodel_cmd_string = "remodel";
    private static final String showmodel_cmd_string = "showmodel";
    private static final String filter_cmd_string = "filter";
    private static final String misassemble_cmd_string = "misassemble";
    private static final Set<String> cmdSet2Disp = new TreeSet();
    static{
        cmdSet2Disp.add(reeval_cmd_string);
        cmdSet2Disp.add(align_eval_cmd_string);
        cmdSet2Disp.add(compmodel_cmd_string);
        Utilities.version=DNAQET.class.getPackage().getImplementationVersion();
    }

    
    public static void main(String[] args){
        
        
        String[] argv = new String[]{"reevaluate", "-s", "RESULT","-d", "HEDE"};
        //String[] argv = new String[]{};
        
        JCommander jc = new JCommander();
        boolean writeLogToDest=true;        

        ReevaluateCommand cre = new ReevaluateCommand();
        AlignEvalCommand aec = new AlignEvalCommand();
        ContigGenerateCommand1 gc1 = new ContigGenerateCommand1();
        ContigGenerateCommand2 gc2 = new ContigGenerateCommand2();
        ContigGenerateCommand3 gc3 = new ContigGenerateCommand3();
        ContigGenerateCommand4 gc4 = new ContigGenerateCommand4();
        ContigGenerateCommand5 gc5 = new ContigGenerateCommand5();
        ComputeModelCommand cmc = new ComputeModelCommand();
        ChrPickerCommand cpc = new ChrPickerCommand();
        RecomputeModelCommand rcmc = new RecomputeModelCommand();
        ShowModelCommand smc = new ShowModelCommand();
        FilterCommand flc = new FilterCommand();
        MisassembleCommand mac = new MisassembleCommand();
                
        
        jc.addCommand(align_eval_cmd_string, aec);
        jc.addCommand(reeval_cmd_string, cre);
        jc.addCommand(compmodel_cmd_string, cmc);        
        jc.addCommand(generate_cmd_string1, gc1);
        jc.addCommand(generate_cmd_string2, gc2);
        jc.addCommand(generate_cmd_string3, gc3);
        jc.addCommand(generate_cmd_string4, gc4);
        jc.addCommand(generate_cmd_string5, gc5);
        jc.addCommand(recompmodel_cmd_string, rcmc);
        jc.addCommand(showmodel_cmd_string, smc);
        jc.addCommand(chrpick_cmd_string, cpc);
        jc.addCommand(filter_cmd_string, flc);
        jc.addCommand(misassemble_cmd_string, mac);
        
        String jarDir="";
        try{
            CodeSource codeSource = DNAQET.class.getProtectionDomain().getCodeSource();
            File jarFile = new File(codeSource.getLocation().toURI().getPath());
            jarDir = jarFile.getParentFile().getPath();            
            jc.parse(args);
//            jc.parse(argv);
            if(jc.getParsedCommand()==null)
                throw new ParameterException("");
            else if(jc.getParsedCommand().equalsIgnoreCase(chrpick_cmd_string)){
                Utilities.current_cmd= cpc;                
                Utilities.current_cmd.startMemoryMeasurement();
                ChromosomePicker cp = new ChromosomePicker(cpc.getInputFile(), cpc.getOutFile(), cpc.getChrFile());
                cp.pickChrs();
                Utilities.current_cmd.endMemoryMeasurement();
            }
            else if(jc.getParsedCommand().equalsIgnoreCase(generate_cmd_string1)){
                Utilities.current_cmd= gc1;                
                Utilities.current_cmd.startMemoryMeasurement();
                ContigGenerator1 g = new ContigGenerator1(gc1.getInputFile(), gc1.getOutputFile(), gc1.getInputContigCoordsFile());
                g.generateContigs();
                Utilities.current_cmd.endMemoryMeasurement();
            }
            else if(jc.getParsedCommand().equalsIgnoreCase(generate_cmd_string2)){
                Utilities.current_cmd= gc2;                
                Utilities.current_cmd.startMemoryMeasurement();
                Utilities.contigsOverlap = gc2.doesContigsOverlap();
                Utilities.nratio = gc2.getNratio();
                ContigGenerator2 g = new ContigGenerator2(gc2.getInputFile(), gc2.getOutputFile(), gc2.getInputContigLengthFile());
                g.generateContigs();
                Utilities.current_cmd.endMemoryMeasurement();
            }
            else if(jc.getParsedCommand().equalsIgnoreCase(generate_cmd_string3)){
                Utilities.current_cmd= gc3;                
                Utilities.current_cmd.startMemoryMeasurement();
                Utilities.contigsOverlap = gc3.doesContigsOverlap();
                Utilities.nratio = gc3.getNratio();
                ContigGenerator3 g = new ContigGenerator3(gc3.getInputFile(), gc3.getOutputFile(), gc3.getMeanContigSize(), gc3.getContigSizeDev(), gc3.getDistribution(), gc3.getNumberFile());
                g.generateContigs();
                Utilities.current_cmd.endMemoryMeasurement();
            }            
            else if(jc.getParsedCommand().equalsIgnoreCase(generate_cmd_string4)){                
                Utilities.current_cmd= gc4;                    
                Utilities.current_cmd.startMemoryMeasurement();
                Utilities.nratio = gc4.getNratio();
                Utilities.depthCoverage = gc4.getCoverage();
                ContigGenerator4 g = new ContigGenerator4(gc4.getInputFile(), gc4.getOutputFile(), gc4.getMeanContigSize(), gc4.getContigSizeDev(), gc4.getDistribution(), gc4.getIsRandom());                                                    
                g.generateContigs();
                Utilities.current_cmd.endMemoryMeasurement();
            }                        
            else if(jc.getParsedCommand().equalsIgnoreCase(generate_cmd_string5)){                
                Utilities.current_cmd= gc5;                    
                Utilities.current_cmd.startMemoryMeasurement();
                Utilities.nratio = gc5.getNratio();
                Utilities.depthCoverage = gc5.getCoverage();
                ContigGenerator5 g = new ContigGenerator5(gc5.getInputFile(), gc5.getOutputFile());                                                    
                g.generateContigs();
                Utilities.current_cmd.endMemoryMeasurement();
            }                                    
            else if(jc.getParsedCommand().equalsIgnoreCase(align_eval_cmd_string)){                
                Utilities.current_cmd = aec;
                Utilities.current_cmd.startMemoryMeasurement();
                boolean modelcomputed=false;
                Utilities.alignmentTool = aec.getAlignmentTool();
                
                // When minimap2 is used the total number of returned alignments are mostly small compared to nucmer,
                // thus use a lower threshold to filter our bad quality alignments                
                if(Utilities.alignmentTool==Enums.AlignmentTool.minimap2)
                        Utilities.alignmentIdentityThreshold=1;
                
                Utilities.no_ReaderThreads = 1;
                Utilities.no_filterThreads = Math.max(1, aec.getThreadNo()-Utilities.no_ReaderThreads);                
                Utilities.resetQSize();
                Utilities.alpha = aec.getAlpha();
                Utilities.beta = aec.getBeta();                
                Utilities.rootDir = aec.getEvalDestFolder();                
                Utilities.submitJobs = aec.getJobCommand();
                Utilities.submitJobsParameters = aec.getJobCommandParams();
                Utilities.no_threads = aec.getThreadNo();
                if(aec.getEvalDestFolder().exists()){
                    Utilities.deleteFolder(aec.getEvalDestFolder());                    
                }            
                // create the sub directories
                File tmpdir;
                File bindir=null;
                for(String d: Utilities.sub_dirs){
                    tmpdir = new File(aec.getEvalDestFolder().getAbsolutePath()+File.separator+d);
                    tmpdir.mkdirs();
                    if(d.equalsIgnoreCase("bin"))
                        bindir = tmpdir;
                }                
                File refFile = aec.getRefFile();
                File divideDestFolder = new File(aec.getEvalDestFolder().getAbsolutePath()+File.separator+"CONTIG_ALIGNMENTS");
                ReferenceSet chrset = new ReferenceSet(refFile, aec.getEvalDestFolder());
                File evalDestFolder = aec.getEvalDestFolder();                                                                
                File alignerDir = new File(jarDir+File.separator+Utilities.alignmentTool.getFolder());
                Utilities.checkAligner(alignerDir);                
                RegressionModel rm=null;                
                if(aec.getModelFile()==null){
                    // create the model file
                    Utilities.assignQuality =false;

                    ModelCompute mc = new ModelCompute(aec.getEvalDestFolder(), refFile, chrset);
                    //modelcomputed=true;
                    rm = mc.process(alignerDir, true);
                    Utilities.assignQuality =true;
                    //Utilities.alignmentIdentityThreshold=tmp;
                }
                else{
                    // load the already existing model file
                    try(ObjectInputStream instream = new ObjectInputStream(new BufferedInputStream(new FileInputStream(aec.getModelFile())));){
                        rm = (RegressionModel)instream.readObject();
                    }
                    catch(ClassNotFoundException e){e.printStackTrace();}
                    catch (IOException ex) {
                        if(!(ex instanceof EOFException))
                            ex.printStackTrace();
                    }                                        
                    if(rm.getReferenceSet()==null || !rm.getReferenceSet().equals(chrset)){
//                    if(rm.getReferenceSet()==null || !rm.getReferenceSet().contains(chrset)){
                        System.out.println("Error: the reference genome used to generate the model doesn't match the reference genome!");
                        System.out.println("Do you want to create a model for this reference and continue the evaluation? (Y/N)");
                        Scanner scan = new Scanner(System.in);
                        String next;
                        while(true){
                            next = scan.next();
                            if(next.equalsIgnoreCase("y") || next.equalsIgnoreCase("yes")){
                                Utilities.assignQuality =false;
                                ModelCompute mc = new ModelCompute(aec.getEvalDestFolder(), refFile, chrset);
                                rm = mc.process(alignerDir, true);  
                                modelcomputed=true;
                                Utilities.assignQuality =true;
                                break;
                            }
                            else if(next.equalsIgnoreCase("n") || next.equalsIgnoreCase("no")){
                                System.out.println("Exiting! Please rerun this command with the correct model file");
                                System.exit(1);                                
                            }
                        }                        
                    }
                    if(!modelcomputed)
                        // Create a copy of the model file in the new evaluation directory
                        Files.copy(aec.getModelFile().toPath(), new File(bindir.getAbsolutePath()+File.separator+"model.bin").toPath(), StandardCopyOption.REPLACE_EXISTING);                    
                }
                if(aec.getLengthScalingFactor()>0)
                    Utilities.lengthScalingFactor = aec.getLengthScalingFactor();
                else
                    Utilities.lengthScalingFactor=rm.getReferenceSet().getMinContigSize();
                
                Utilities.regm = rm ;
                
                divideDestFolder.mkdirs();                                                                                                                
                int numberofparts = aec.getNumberPartitions();                
                File inputfile = aec.getInputFile();
                Utilities.current_cmd.startTimer4Process("Aligning the assembly back to reference");
                DataAligner d = new DataAligner(inputfile, numberofparts, divideDestFolder, alignerDir, chrset);            
                //int n=d.setAligners(Utilities.current_cmd);
                int n=d.setAligners(false);
                if(n<0)
                    return;
                Utilities.waiter(d.getTouchFiles());                
                Utilities.current_cmd.endTimer4Process("Aligning the assembly back to reference");                                
                Utilities.alignmentDistanceThreshold= aec.getDistThreshold();                
                
                // Read the alignment files
                Utilities.current_cmd.startTimer4Process("Reading the alignments and computing quality scores");
                Evaluator ev = new Evaluator(divideDestFolder, evalDestFolder, chrset);
                ev.processFiles();
                Utilities.current_cmd.endTimer4Process("Reading the alignments and computing quality scores");      
                Utilities.deleteFolder(divideDestFolder);
                Utilities.deleteFolder(chrset.getRefFolder());
                Utilities.current_cmd.endMemoryMeasurement();
            }
            else if(jc.getParsedCommand().equalsIgnoreCase(reeval_cmd_string)){                
                RegressionModel rm=null;
                Utilities.current_cmd = cre;
                Utilities.current_cmd.startMemoryMeasurement();
                
                Utilities.alpha = cre.getAlpha();
                Utilities.beta = cre.getBeta();
                
                File evalDestFolder = cre.getEvalDestFolder();
                Utilities.rootDir = evalDestFolder;
                File sourceFolder = cre.getSourceFolder(); 
                Utilities.no_ReaderThreads = 1;
                Utilities.no_filterThreads = Math.max(1, cre.getThreadNo()-Utilities.no_ReaderThreads);
                Utilities.resetQSize();

                Utilities.alignmentDistanceThreshold= cre.getDistThreshold();    
                ReferenceSet chrset = ReferenceSet.load(sourceFolder);
                File modfile=null;
                if(modfile==null){
                    // use the already existing model in the previously evaluated directory
                    modfile = new File(cre.getSourceFolder()+File.separator+Utilities.modelFileName);
                }
                                
                try(ObjectInputStream instream = new ObjectInputStream(new BufferedInputStream(new FileInputStream(modfile)));){
                    rm = (RegressionModel)instream.readObject();
                }
                catch(ClassNotFoundException e){e.printStackTrace();}
                catch (IOException ex) {
                    if(!(ex instanceof EOFException))
                        ex.printStackTrace();
                }
                Utilities.regm = rm ;                
                if(cre.getLengthScalingFactor()>0)
                    Utilities.lengthScalingFactor = cre.getLengthScalingFactor();
                else
                    Utilities.lengthScalingFactor = rm.getReferenceSet().getMinContigSize();                
                                                               
                Utilities.current_cmd.startTimer4Process("Reading the alignments and computing quality scores");
                Reevaluator ev = new Reevaluator(sourceFolder, evalDestFolder, chrset, modfile);
                writeLogToDest=ev.processFiles();
                Utilities.current_cmd.endTimer4Process("Reading the alignments and computing quality scores");                                
                Utilities.current_cmd.endMemoryMeasurement();
            }
            else if(jc.getParsedCommand().equalsIgnoreCase(compmodel_cmd_string)){
                // now we are serious
                Utilities.current_cmd = cmc;
                Utilities.current_cmd.startMemoryMeasurement();
                Utilities.assignQuality =false;                
                Utilities.alignmentTool = cmc.getAlignmentTool();

                // When minimap2 is used the total number of returned alignments are mostly small compared to nucmer,
                // thus use a lower threshold to filter our bad quality alignments                
                if(Utilities.alignmentTool==Enums.AlignmentTool.minimap2)
                        Utilities.alignmentIdentityThreshold=1;                
                Utilities.submitJobs = cmc.getJobCommand();
                Utilities.submitJobsParameters = cmc.getJobCommandParams();
                Utilities.no_threads = cmc.getThreadNo();                
                Utilities.no_ReaderThreads = 1;
                Utilities.no_filterThreads = Math.max(1, cmc.getThreadNo()-Utilities.no_ReaderThreads);                
                Utilities.resetQSize();
                Utilities.rootDir = cmc.getDestinationFolder();

                File alignerDir = new File(jarDir+File.separator+Utilities.alignmentTool.getFolder());
                Utilities.checkAligner(alignerDir);                
                
                if(!cmc.getDestinationFolder().exists())
                    cmc.getDestinationFolder().mkdirs();
                File bindir = new File(cmc.getDestinationFolder().getAbsolutePath()+File.separator+"bin");
                bindir.mkdirs();
                ReferenceSet chrset = new ReferenceSet(cmc.getInputFile(), cmc.getDestinationFolder());
                ModelCompute mc = new ModelCompute(cmc.getDestinationFolder(), cmc.getInputFile(), chrset);
                mc.process(alignerDir, true);
                File modelfile = mc.getModelFile();
                Utilities.deleteFolder(chrset.getRefFolder());
                File targetfile = new File(cmc.getDestinationFolder().getAbsolutePath()+File.separator+"model.bin");
                Files.move(modelfile.toPath(), targetfile.toPath(), StandardCopyOption.REPLACE_EXISTING);                
                Utilities.deleteFolder(bindir);
                Utilities.current_cmd.endMemoryMeasurement();
            }
            else if(jc.getParsedCommand().equalsIgnoreCase(recompmodel_cmd_string)){
                Utilities.current_cmd = rcmc;
                Utilities.current_cmd.startMemoryMeasurement();
                Reevaluator4Artifact r4a = new Reevaluator4Artifact(rcmc.getInputFile(), rcmc.getOutFile());
                r4a.process();
                String toprint =ShowModel.getModelDetails(rcmc.getOutFile());
                Utilities.current_cmd.endMemoryMeasurement();
                System.out.println(toprint);                
            }
            else if(jc.getParsedCommand().equalsIgnoreCase(showmodel_cmd_string)){
                Utilities.current_cmd = smc;
                Utilities.current_cmd.startMemoryMeasurement();
                Utilities.showModelData = smc.getShowModelData();
                Utilities.showModelInTabularFormat = smc.isIntabularformat();
                String toprint = ShowModel.getModelDetails(smc.getInputFile());
                Utilities.current_cmd.endMemoryMeasurement();
                System.out.println(toprint);                
            }
            else if(jc.getParsedCommand().equalsIgnoreCase(filter_cmd_string)){
                Utilities.fastarecordlengththreshold = flc.getLength();
                Utilities.current_cmd = flc;
                Utilities.current_cmd.startMemoryMeasurement();
                FastaFilter f = new FastaFilter(flc.getInputFile(), flc.getOutputFile());
                f.filter();
                Utilities.current_cmd.endMemoryMeasurement();
                
            }
            else if(jc.getParsedCommand().equalsIgnoreCase(misassemble_cmd_string)){
                Utilities.current_cmd = mac;
                File infile = mac.getInputFile();
                File outfile = mac.getOutputFile();
                int misThre = mac.getMis_thre();
                int number = mac.getAvgmisassembly();
                                
                Utilities.current_cmd.startMemoryMeasurement();
                MisassemblyGenerator mg = new MisassemblyGenerator(infile, outfile, misThre, number);
                mg.process();
                //FastaFilter f = new FastaFilter(flc.getInputFile(), flc.getOutputFile());
                //f.filter();
                Utilities.current_cmd.endMemoryMeasurement();
                
            }
            
            if(Utilities.current_cmd!=null && Utilities.current_cmd!=smc){
                Utilities.current_cmd.writeCommand(writeLogToDest);
            }
        }
        catch(ParameterException pr){
            System.out.println("dnAQET version "+Utilities.version);
            System.out.println(pr.getMessage());
            if(args.length>0){
                JCommander com=jc.getCommands().get(args[0]);
                if(com==null){                    
                    //jc.usage();
                    printUsage(jc);
                }
                else{
                    com.usage();                
                }
            }
            else{
                //jc.usage();
                printUsage(jc);
            }
        }
        catch(ExceptionInInitializerError ex){
            ex.printStackTrace();
        } catch (IOException ex) {
            ex.printStackTrace();
        } catch (URISyntaxException ex) {
            ex.printStackTrace();
        }
        
    }
    private static void printUsage(JCommander jc){
        
        System.out.println("Usage: <main class> [options] [command] [command options]"+Utilities.ls);
        for(String s : cmdSet2Disp){
            JCommander com=jc.getCommands().get(s);
            System.out.println("Command: "+s+"\t"+jc.getCommandDescription(s));
            com.usage();
        }
    }
}
