// -*- LSST-C++ -*-

#include <cmath>
#include <fstream>
#include <iostream>
#include <sstream>

#include <unistd.h>
#include <getopt.h>

#include "collapseTracklets.h"
#include "rmsLineFit.h"


namespace rmsLineFit {





    double getAverageMagnitude(const Tracklet t, const::std::vector<Detection>* detections) {
        double sum = 0.;
        unsigned int count = 0;
        std::set<unsigned int>::const_iterator detIndexIter;
        for (detIndexIter = t.indices.begin(); detIndexIter != t.indices.end(); detIndexIter++) {
            sum += (*detections)[*detIndexIter].getMag();
            count++;
        }
        if (count < 1) {
            // TBD: real LSST exceptions
            std::cerr << "EE: programming error: getAverageMagnitude: highly unexpected error - tracklet has no detections?" 
                      << std::endl;
            exit(1);
        }
        return sum/count;
    }







    double rmsForTracklet(const Tracklet t, const std::vector<Detection> *detections, 
                          std::vector<double> *perDetSqDist) {
        std::vector<double> RASlopeAndOffset, DecSlopeAndOffset;
        std::set<unsigned int>::iterator indicesIter;
        std::vector<Detection> trackletDets;
        std::vector<Detection>::iterator detIter;
        for (indicesIter = t.indices.begin(); 
             indicesIter != t.indices.end();
             indicesIter++) {
            trackletDets.push_back((*detections)[*indicesIter]);
        }
        
        if ((perDetSqDist != NULL) && (perDetSqDist->size() != 0)) {
            // TBD: real exceptions of course
            std::cerr << "EE: PROGRAMMING ERROR: perDetSqDist argument to rmsForTracklet must be either NULL or a pointer to an allocated, empty vector!\n";
            exit(1);
        }
        

        if (t.indices.size() == 0) {
            std:: cerr << "EE: PROGRAMMING ERROR: rmsForTracklet received a tracklet associated with 0 detection indices.\n";
            exit(1);
        }
        
        /* set the first detection time as t = 0 (timeOffset) */
        double t0 = trackletDets[0].getEpochMJD();
        collapseTracklets::TrackletCollapser myTC;
        myTC.leastSquaresSolveForRADecLinear(&trackletDets, RASlopeAndOffset, 
                                                           DecSlopeAndOffset, t0);
        /* since our current least squares library doesn't give it to us, we have to re-calculate
         * what that least square was...*/
        double squaresSum = 0.0;
        for (detIter = trackletDets.begin(); detIter != trackletDets.end(); detIter++) {
            double tOffset = detIter->getEpochMJD() - t0;

            double projectedRA = RASlopeAndOffset[0] * tOffset + RASlopeAndOffset[1];
	    projectedRA = KDTree::Common::convertToStandardDegrees(projectedRA);

            double projectedDec = DecSlopeAndOffset[0] * tOffset + DecSlopeAndOffset[1];
	    projectedDec = KDTree::Common::convertToStandardDegrees(projectedDec);

	    double actualRA = KDTree::Common::convertToStandardDegrees(detIter->getRA());
	    double actualDec = KDTree::Common::convertToStandardDegrees(detIter->getDec());

	    // 	    std::cout << " projected position: (" << projectedRA << ", " << projectedDec << ")" << std::endl;
	    // 	    std::cout << " Actual position: (" << actualRA << ", " << actualDec <<  ")" << std::endl;

            double RADist = fabs(projectedRA -  actualRA);
            double DecDist = fabs(projectedDec -  actualDec);
	    if (RADist > 180.) {
	      RADist = fabs(RADist - 360);
	    }
	    if (DecDist > 180.) {
	      DecDist = fabs(DecDist - 360);
	    }
	    if ((RADist > 180.)  || (DecDist > 180.)) {
		  // TBD: LSST exceptions
                std::cerr <<  "EE: rmsLineFit: Unexpected programming error\n";
                exit(1);
            }
	    //	    std::cout << " got radist, decdist = " << RADist << ", " << DecDist << std::endl;
	    double localDistSquared = RADist*RADist + DecDist*DecDist;
	    //	    std::cout << " local distsquared = " << localDistSquared << std::endl;
            /* get the sum of all the squares of distances */
            if (perDetSqDist != NULL) {
                perDetSqDist->push_back(localDistSquared);
            }
            squaresSum += localDistSquared;
        }
	double result = sqrt(squaresSum) / trackletDets.size(); 
	return  result;
    }



    


    
    void filterByLineFitAddToOutputVector(const std::vector<Tracklet> *tracklets, 
                                          const std::vector<Detection> *allDets,
                                          double maxRMSm, double maxRMSb,
                                          std::vector<Tracklet> &output) {
        std::vector<Tracklet>::const_iterator tIter;
        for (tIter = tracklets->begin(); tIter != tracklets->end(); tIter++) {
	  double rms = rmsForTracklet(*tIter, allDets);
	  // std::cout << "got RMS = " << rms << ", tracklet length = " <<  tIter->indices.size() << std::endl;
	  double maxRMS =  maxRMSm * getAverageMagnitude(*tIter, allDets) + maxRMSb;
	  if ( maxRMS >= rms) {
                output.push_back(*tIter);
            }
        }        
    }




    
    /*
     * given a pair of functions relating MJD to RA, Dec s.t. RA = RASlope*(MJD - tOffset) + RAintercept and
     * the similarly for declination, return a map which relates indices of *t to squared distance to the projected point.
     */
    std::map <unsigned int, double> getPerDetSqDistanceToLine(const Tracklet *t, const std::vector<Detection>* allDets, 
                                                              double RASlope, double RAIntercept, double DecSlope, double DecIntercept, 
                                                              double tOffset) {
        std::map <unsigned int, double> results;
        for (std::set<unsigned int>::const_iterator indicesIter = t->indices.begin();
             indicesIter != t->indices.end();
             indicesIter++) {
            double t = (*allDets)[*indicesIter].getEpochMJD() - tOffset;
            double projectedRA = KDTree::Common::convertToStandardDegrees(RASlope * t + RAIntercept);
            double projectedDec = KDTree::Common::convertToStandardDegrees(DecSlope *t + DecIntercept);
            double obsRA = KDTree::Common::convertToStandardDegrees((*allDets)[*indicesIter].getRA());
            double obsDec = KDTree::Common::convertToStandardDegrees((*allDets)[*indicesIter].getDec());
            double RADist = fabs(projectedRA - obsRA);
            if (RADist > 180.) {
                RADist = fabs(RADist - 360);
                if (RADist > 180.) {
                    std::cerr <<  "EE: getPerDetSqDistanceToLine: Unexpected programming error: failed to get distance < 360 between two degree points\n";
                    exit(1);
                }
            }
            double DecDist = fabs(projectedDec - obsDec);
            if (DecDist > 180.) {
                DecDist = fabs(DecDist - 360);
                if (DecDist > 180.) {
                    std::cerr <<  "EE: Unexpected programming error: failed to get distance < 360 between two degree points\n";
                    exit(1);
                }
            }
            double sqDist = RADist*RADist + DecDist*DecDist;
            results.insert(std::make_pair(*indicesIter, sqDist));            
        }
        return results;
    }



    std::vector<Detection> getTrackletDets(const Tracklet *t, const std::vector<Detection>* allDets) {
        std::vector<Detection> results;
        for (std::set<unsigned int>::const_iterator iIter = t->indices.begin();
             iIter != t->indices.end();
             iIter++) {
            results.push_back((*allDets)[*iIter]);
        }
        return results;
    }



    
    Tracklet TrackletPurifier::purifyTracklet(const Tracklet *t, const std::vector<Detection>* allDets, 
                                              double maxRMSm, double maxRMSb) {
        Tracklet curTracklet = *t;
        bool isClean = false;

        while (isClean == false) {
            double t0 = (*allDets)[*(curTracklet.indices.begin())].getEpochMJD();
            std::vector<Detection> curTrackletDets = getTrackletDets(&curTracklet, allDets);
            std::vector<double> RASlopeAndOffset, DecSlopeAndOffset;
            collapseTracklets::TrackletCollapser myTC;
            myTC.leastSquaresSolveForRADecLinear(&curTrackletDets, RASlopeAndOffset, 
                                                 DecSlopeAndOffset, t0);
            std::map<unsigned int, double> indexToSqDist = 
                getPerDetSqDistanceToLine(&curTracklet, allDets, RASlopeAndOffset[0], RASlopeAndOffset[1],
                                          DecSlopeAndOffset[0], DecSlopeAndOffset[1], t0);

            isClean = true;
            double worstDetVal = 0.0;
            unsigned int worstDetIndex = 0;
            for (std::map<unsigned int, double>::iterator distIter = indexToSqDist.begin();
                 distIter != indexToSqDist.end(); distIter++) {
                double distMax = maxRMSm * (*allDets)[distIter->first].getMag() + maxRMSb;
                if ((distIter->second > distMax*distMax) && (distIter->second > worstDetVal)) {
                    worstDetVal = distIter->second;
		    if (worstDetVal > 1) {
                        std::cerr << "Warning: detection point to projected point is improbably large distance: " << worstDetVal << std::endl;
		    }
                    worstDetIndex = distIter->first;
                    isClean = false;
                }
            }
            if (isClean == false) {
                curTracklet.indices.erase(worstDetIndex);
            }
        }
        return curTracklet;
    }






    int rmsLineFitMain(int argc, char** argv) {
        std::string USAGE("USAGE: lineFitFilter --detsFile <detections file> --pairsFile <tracklets (pairs) file) [--maxRMSm <number> --maxRMSb <number> ] --outFile <output tracklets (pairs) file>");
        char* pairsFileName = NULL;
        char* detsFileName = NULL;
        char* outFileName = NULL;

        std::vector <Tracklet> trackletsVector;
        std::vector <Detection> detsVector;
        
        double maxRMSb = .001;
        double maxRMSm = 0.;

        static const struct option longOpts[] = {
            { "pairsFile", required_argument, NULL, 'p' },
            { "detsFile", required_argument, NULL, 'd' },
            { "outFile", required_argument, NULL, 'o' },
            { "maxRMSm", required_argument, NULL, 'm'},
            { "maxRMSb", required_argument, NULL, 'b'},
            { "help", no_argument, NULL, 'h' },
            { NULL, no_argument, NULL, 0 }
        };

        int longIndex = -1;
        const char* optString = "p:d:o:m:b:h";        
        int opt = getopt_long( argc, argv, optString, longOpts, &longIndex );        
        std::stringstream ss;

        while( opt != -1 ) {
            switch( opt ) {
            case 'p':
                pairsFileName = optarg; 
                break;
                
            case 'd':
                detsFileName = optarg ; 
                break;
                
            case 'o':
                outFileName = optarg;
                break;
                
            case 'm':
                ss.clear();
                ss << optarg;
                ss >> maxRMSm;
                break;

            case 'b':
                ss.clear();
                ss << optarg;
                ss >> maxRMSb;
                break;
                
            case 'h':   /* fall-through is intentional */
            case '?':
                std::cout << "got request for help " << std::endl;
                std::cout<<USAGE<<std::endl;
                return 0;
                break;
            default:
                std::cerr <<  "EE: Unexpected programmer error in options parsing\n";
                exit(1);
                break;
            }        
            opt = getopt_long( argc, argv, optString, longOpts, &longIndex );
        }

        if ((pairsFileName == NULL) || (detsFileName == NULL) || (outFileName == NULL)) {
            std::cerr << "Did not get required parameters " << std::endl;
            std::cerr << USAGE << std::endl;
            return 1;
        }
        std::ifstream pairsFile(pairsFileName);
        std::ifstream detsFile(detsFileName);
        std::ofstream outFile(outFileName);
        if (!detsFile.is_open()) {
            std::cerr << "EE: Failed to open dets file " << detsFileName << " - does this file exist?" 
                      << std::endl;
            return 1;
        }
        if (!pairsFile.is_open()) {
            std::cerr << "EE: Failed to open pairs file " << pairsFileName << " - does this file exist?" 
                      << std::endl;
            return 1;
        }
        if (!outFile.is_open()) {
            std::cerr << "EE: Failed to open output file " << outFileName << " - do you have write permissions?" 
                      << std::endl;
            return 1;
        }
        

        std::cout << "lineFitFilter: " << std::endl;
        std::cout << "==========================================" << std::endl;
        std::cout << "Detections file:        " << detsFileName << std::endl;
        std::cout << "Pairs (tracklets) file: " << pairsFileName << std::endl;
        std::cout << "Output file:            " << outFileName << std::endl;
        std::cout << "max RMS:                " 
                  << maxRMSm << " * av. magnitude of tracklet + " << maxRMSb << std::endl;

        std::cout << "Reading detections file...." << std::endl;
        collapseTracklets::TrackletCollapser myTC;
        myTC.populateDetVectorFromFile(detsFile, detsVector);
        std::cout << "Reading tracklets (pairs) file...." << std::endl;
        myTC.populatePairsVectorFromFile(pairsFile, trackletsVector);
        std::cout << "Done!" << std::endl;

        if (!myTC.isSane(detsVector.size(), &trackletsVector)) {
            std::cerr << "EE: Pairs file does not seem to correspond with detections file.  Exiting." 
                      << std::endl;
            return 3;
        }


        std::vector<Tracklet> postFilteredTracklets;
        std::cout << "Doing the filtering..." << std::endl;
        filterByLineFitAddToOutputVector(&trackletsVector, &detsVector, maxRMSm, maxRMSb, postFilteredTracklets);
        std::cout << "Done. Writing output." << std::endl;
        myTC.writeTrackletsToOutFile(&postFilteredTracklets, outFile);

        return 0;
    }



    int rmsPurifyMain(int argc, char** argv) {
        std::string USAGE("USAGE: purifyTracklets --detsFile <detections file> --pairsFile <tracklets (pairs) file) [--maxRMSm <number> --maxRMSb <number> ] --outFile <output tracklets (pairs) file>");
        char* pairsFileName = NULL;
        char* detsFileName = NULL;
        char* outFileName = NULL;

        std::vector <Tracklet> trackletsVector;
        std::vector <Detection> detsVector;
        
        double maxRMSb = .001;
        double maxRMSm = 0.;
        unsigned int minObs = 2;

        static const struct option longOpts[] = {
            { "pairsFile", required_argument, NULL, 'p' },
            { "detsFile", required_argument, NULL, 'd' },
            { "outFile", required_argument, NULL, 'o' },
            { "maxRMSm", required_argument, NULL, 'm'},
            { "maxRMSb", required_argument, NULL, 'b'},
            { "minobs", required_argument, NULL, 'n'},
            { "help", no_argument, NULL, 'h' },
            { NULL, no_argument, NULL, 0 }
        };

        int longIndex = -1;
        const char* optString = "p:d:o:m:b:h";        
        int opt = getopt_long( argc, argv, optString, longOpts, &longIndex );        
        std::stringstream ss;

        while( opt != -1 ) {
            switch( opt ) {
            case 'p':
                pairsFileName = optarg; 
                break;
                
            case 'd':
                detsFileName = optarg ; 
                break;
                
            case 'o':
                outFileName = optarg;
                break;
                
            case 'm':
                ss.clear();
                ss << optarg;
                ss >> maxRMSm;
                break;

            case 'n':
                ss.clear();
                ss << optarg;
                ss >> minObs;
                break;

            case 'b':
                ss.clear();
                ss << optarg;
                ss >> maxRMSb;
                break;
                
            case 'h':   /* fall-through is intentional */
            case '?':
                std::cout << "got request for help " << std::endl;
                std::cout<<USAGE<<std::endl;
                return 0;
                break;
            default:
                std::cerr << "EE: Unexpected programmer error in options parsing\n";
                exit(1);
                break;
            }        
            opt = getopt_long( argc, argv, optString, longOpts, &longIndex );
        }

        if ((pairsFileName == NULL) || (detsFileName == NULL) || (outFileName == NULL)) {
            std::cerr << "Did not get required parameters " << std::endl;
            std::cerr << USAGE << std::endl;
            return 1;
        }
        std::ifstream pairsFile(pairsFileName);
        std::ifstream detsFile(detsFileName);
        std::ofstream outFile(outFileName);
        if (!detsFile.is_open()) {
            std::cerr << "EE: Failed to open dets file " << detsFileName << " - does this file exist?" 
                      << std::endl;
            return 1;
        }
        if (!pairsFile.is_open()) {
            std::cerr << "EE: Failed to open pairs file " << pairsFileName << " - does this file exist?" 
                      << std::endl;
            return 1;
        }
        if (!outFile.is_open()) {
            std::cerr << "EE: Failed to open output file " << outFileName << " - do you have write permissions?" 
                      << std::endl;
            return 1;
        }
        
        std::cout << "purifyTracklets: " << std::endl;
        std::cout << "==========================================" << std::endl;
        std::cout << "Detections file:        " << detsFileName << std::endl;
        std::cout << "Pairs (tracklets) file: " << pairsFileName << std::endl;
        std::cout << "Output file:            " << outFileName << std::endl;
        std::cout << "max RMS:                " 
                  << maxRMSm << " * av. magnitude of tracklet + " << maxRMSb << std::endl;

        std::cout << "Reading detections file...." << std::endl;
        collapseTracklets::TrackletCollapser myTC;
        myTC.populateDetVectorFromFile(detsFile, detsVector);
        std::cout << "Reading tracklets (pairs) file...." << std::endl;
        myTC.populatePairsVectorFromFile(pairsFile, trackletsVector);
        std::cout << "Done!" << std::endl;

        if (!myTC.isSane(detsVector.size(), &trackletsVector)) {
            std::cerr << "EE: Pairs file does not seem to correspond with detections file.  Exiting." 
                      << std::endl;
            return 3;
        }

        std::vector<Tracklet> postFilteredTracklets;
        std::cout << "Doing the filtering..." << std::endl;
        std::vector<Tracklet>::iterator tIter;
        TrackletPurifier myTP;
        for (tIter = trackletsVector.begin(); tIter != trackletsVector.end(); tIter++) {
            Tracklet tmp = myTP.purifyTracklet(&(*tIter), &detsVector, maxRMSm, maxRMSb);
            if (tmp.indices.size() >= minObs) {
                postFilteredTracklets.push_back(tmp);
            }
        }        
        std::cout << "Done. Writing output." << std::endl;
        myTC.writeTrackletsToOutFile(&postFilteredTracklets, outFile);

        return 0;
    }


}
