#include <animaDTITractographyImageFilter.h>
#include <itkTimeProbe.h>

#include <animaReadWriteFunctions.h>
#include <animaLogTensorImageFilter.h>
#include <itkCommand.h>

#include <tclap/CmdLine.h>

#include <animaShapesWriter.h>

//Update progression of the process
void eventCallback (itk::Object* caller, const itk::EventObject& event, void* clientData)
{
    itk::ProcessObject * processObject = (itk::ProcessObject*) caller;
    std::cout<<"\033[K\rProgression: "<<(int)(processObject->GetProgress() * 100)<<"%"<<std::flush;
}

int main(int argc,  char*  argv[])
{
    TCLAP::CmdLine cmd("INRIA / IRISA - VisAGeS/Empenn Team", ' ',ANIMA_VERSION);

    TCLAP::ValueArg<std::string> dtiArg("i","dti","Input DT image",true,"","DT image",cmd);
    TCLAP::ValueArg<std::string> seedMaskArg("s","seed-mask","Seed mask",true,"","seed",cmd);
    TCLAP::ValueArg<std::string> fibersArg("o","fibers","Output fibers",true,"","fibers",cmd);

    TCLAP::ValueArg<std::string> filterMaskArg("","filter-mask","Mask for filtering fibers (default: none)",false,"","filter mask",cmd);
    TCLAP::ValueArg<std::string> cutMaskArg("c","cut-mask","Mask for cutting fibers (default: none)",false,"","cut mask",cmd);
    TCLAP::ValueArg<std::string> forbiddenMaskArg("f","forbidden-mask","Mask for removing fibers (default: none)",false,"","remove mask",cmd);

    TCLAP::ValueArg<double> adcThrArg("A","adc-thr","ADC threshold (default: 2.0e-3)",false,2.0e-3,"ADC threshold",cmd);
    TCLAP::ValueArg<double> faThrArg("","fa-thr","FA threshold (default: 0.1)",false,0.1,"FA threshold",cmd);
    TCLAP::ValueArg<double> punctureWeightArg("p","punc-weight","Puncture weight to go through planar tensors (default: 0.2)",false,0.2,"puncture weight",cmd);

    TCLAP::ValueArg<double> stopAngleArg("a","angle-max","Maximum angle for tracking (default: 60)",false,60.0,"maximum track angle",cmd);
    TCLAP::ValueArg<double> stepLengthArg("","step-length","Length of each step (in mm, default: 0.5)",false,0.5,"step length",cmd);
    TCLAP::ValueArg<int> nbFibersArg("","nb-fibers","Number of starting filters (n*n*n) per voxel (default: 1)",false,1,"number of seeds per voxel",cmd);

    TCLAP::ValueArg<double> minLengthArg("","min-length","Minimum length for a fiber to be considered for computation (default: 10mm)",false,10.0,"minimum length",cmd);
    TCLAP::ValueArg<double> maxLengthArg("","max-length","Maximum length of a tract (default: 200mm)",false,200.0,"maximum length",cmd);

    TCLAP::SwitchArg addLocalDataArg("L","local-data","Add local data information to output tracks",cmd);

    TCLAP::ValueArg<unsigned int> nbThreadsArg("T","nb-threads","Number of threads to run on (default: all available)",false,itk::MultiThreaderBase::GetGlobalDefaultNumberOfThreads(),"number of threads",cmd);

    try
    {
        cmd.parse(argc,argv);
    }
    catch (TCLAP::ArgException& e)
    {
        std::cerr << "Error: " << e.error() << "for argument " << e.argId() << std::endl;
        return EXIT_FAILURE;
    }

    typedef anima::dtiTractographyImageFilter MainFilterType;
    typedef MainFilterType::ModelImageType DTIImageType;
    typedef MainFilterType::MaskImageType MaskImageType;

    MainFilterType::Pointer dtiTracker = MainFilterType::New();

    typedef anima::LogTensorImageFilter <DTIImageType::IOPixelType,DTIImageType::ImageDimension> LogTensorFilterType;
    LogTensorFilterType::Pointer tensorLogger = LogTensorFilterType::New();

    tensorLogger->SetInput(anima::readImage <DTIImageType> (dtiArg.getValue()));
    tensorLogger->SetScaleNonDiagonal(false);
    tensorLogger->SetNumberOfWorkUnits(nbThreadsArg.getValue());

    tensorLogger->Update();

    dtiTracker->SetInputImage(tensorLogger->GetOutput());

    dtiTracker->SetSeedingMask(anima::readImage <MaskImageType> (seedMaskArg.getValue()));

    if (filterMaskArg.getValue() != "")
        dtiTracker->SetFilteringMask(anima::readImage <MaskImageType> (filterMaskArg.getValue()));

    if (cutMaskArg.getValue() != "")
        dtiTracker->SetCutMask(anima::readImage <MaskImageType> (cutMaskArg.getValue()));

    if (forbiddenMaskArg.getValue() != "")
        dtiTracker->SetForbiddenMask(anima::readImage <MaskImageType> (forbiddenMaskArg.getValue()));

    dtiTracker->SetNumberOfWorkUnits(nbThreadsArg.getValue());
    dtiTracker->SetNumberOfFibersPerPixel(nbFibersArg.getValue());
    dtiTracker->SetStepProgression(stepLengthArg.getValue());
    dtiTracker->SetStopFAThreshold(faThrArg.getValue());
    dtiTracker->SetStopADCThreshold(adcThrArg.getValue());
    dtiTracker->SetPunctureWeight(punctureWeightArg.getValue());
    dtiTracker->SetMaxFiberAngle(stopAngleArg.getValue());
    dtiTracker->SetMinLengthFiber(minLengthArg.getValue());
    dtiTracker->SetMaxLengthFiber(maxLengthArg.getValue());

    bool computeColors = (fibersArg.getValue().find(".fds") != std::string::npos) && (addLocalDataArg.isSet());
    dtiTracker->SetComputeLocalColors(computeColors);

    itk::CStyleCommand::Pointer callback = itk::CStyleCommand::New();
    callback->SetCallback(eventCallback);
    dtiTracker->AddObserver(itk::ProgressEvent(), callback);

    itk::TimeProbe tmpTime;
    tmpTime.Start();

    dtiTracker->Update();
    dtiTracker->RemoveAllObservers();

    tmpTime.Stop();
    std::cout << "Tracking time: " << tmpTime.GetTotal() << "s" << std::endl;

    anima::ShapesWriter writer;
    writer.SetInputData(dtiTracker->GetOutput());
    writer.SetFileName(fibersArg.getValue());
    writer.Update();

    return EXIT_SUCCESS;
}
