#include <set>
#include <list>
#include <vector>
#include <cmath>
#include <algorithm>
#include <iostream>

#include <TH2F.h>
#include <TMath.h>

#include <HEPUnits.hxx>
#include <ICOMETLog.hxx>

#include "IHoughLines.hxx"

COMET::IHoughPoint::IHoughPoint() : fAngle(0) {}
COMET::IHoughPoint::~IHoughPoint() {}

void COMET::IHoughPoint::RemoveFromLines(double threshold) {
    // Find the maximum weight.
    double maxWeight = -99999.0;
    for (std::set<COMET::IHoughLine*>::iterator ln =  fLines.begin();
         ln != fLines.end();
         ++ln) {
        double weight = (*ln)->GetWeight();
        maxWeight = std::max(maxWeight,weight);
    }
    
    std::set<COMET::IHoughLine*>::iterator ln =  fLines.begin();
    while (ln != fLines.end()) {
        double weight = (*ln)->GetWeight();
        if (weight >= threshold*maxWeight) {
            ++ln;
            continue;
        }
        std::list<COMET::IHoughPoint*>::iterator pt
            = std::find((*ln)->fPoints.begin(),(*ln)->fPoints.end(),this);
        if (pt == (*ln)->fPoints.end()) {
            COMETWarn("RemoveFromLines:"
                      << " OUCH! Point not in line");
            ++ln;
            continue;
        }
        (*ln)->fPoints.erase(pt);
        std::set<COMET::IHoughLine*>::iterator target = ln;
        ++ln;
        fLines.erase(target);
    }
}

void COMET::IHoughPoint::KeepLine(COMET::IHoughLine* line) {
    std::set<COMET::IHoughLine*>::iterator ln
        = std::find(fLines.begin(),fLines.end(),line);
    if (ln == fLines.end()) {
        COMETWarn("COMET::IHoughPoint::KeepLine: Line doesn't contain point");
        return;
    }
    ln = fLines.begin();
    while (ln!=fLines.end()) {
        if (*ln == line) {
            ++ln;
            continue;
        }
        std::list<COMET::IHoughPoint*>::iterator pt
            = std::find((*ln)->fPoints.begin(),(*ln)->fPoints.end(),this);
        if (pt == (*ln)->fPoints.end()) {
            COMETWarn("COMET::IHoughPoint::RemoveFromLines:"
                      << " OUCH! Point not in line");
            ++ln;
            continue;
        }
        (*ln)->fPoints.erase(pt);
        std::set<COMET::IHoughLine*>::iterator target = ln;
        ++ln;
        fLines.erase(target);
    }
}

COMET::IHoughLine::IHoughLine(COMET::IHoughLines* parent, double angle, double radius) 
    : fParent(parent), fAngle(angle), fRadius(radius), fSorted(true) { }

COMET::IHoughLine::~IHoughLine() {}

double COMET::IHoughLine::GetWeight() const {
    double weight = 0.0;
    for (std::list<COMET::IHoughPoint*>::const_iterator p = fPoints.begin();
         p != fPoints.end();
         ++p) {
        weight += (*p)->GetWeight();
    }
    return weight;
}

void COMET::IHoughLine::RemovePoints(Points::iterator begin, Points::iterator end) {
    for (Points::iterator tmp = begin; tmp != end; ++tmp) {
        (*tmp)->fLines.erase(this);
    }
    fPoints.erase(begin,end);
}

void COMET::IHoughLine::AddPoint(COMET::IHoughPoint* point) {
    fPoints.push_back(point);
    point->fLines.insert(this);
    fSorted = false;
}

void COMET::IHoughLine::Clear() {
    fPoints.clear();
}

double COMET::IHoughLine::GetLength(const COMET::IHoughPoint* point) const {
    double dX = point->GetX() - fParent->GetCenterX();
    double dY = point->GetY() - fParent->GetCenterY();
    double radius = std::sqrt(dX*dX + dY*dY);
    // The angle of the line is wrt the X axis, but want the angle to the
    // point of closest approach.  That angle is at 90 deg wrt the angle of
    // the line.  This is done explicitly to make the geometry painfully
    // clear.
    double dTheta = point->GetAngle() - (fAngle + M_PI_2);
    return radius*std::sin(dTheta);
}

double COMET::IHoughLine::GetDeviation(const COMET::IHoughPoint* point) const {
    double dX = point->GetX() - fParent->GetCenterX();
    double dY = point->GetY() - fParent->GetCenterY();
    double radius = std::sqrt(dX*dX + dY*dY);
    // The angle of the line is wrt the X axis, but want the angle to the
    // point of closest approach.  That angle is at 90 deg wrt the angle of
    // the line.  This is done explicitly to make the geometry painfully
    // clear.
    double dTheta = point->GetAngle() - (fAngle + M_PI_2);
    return radius*std::cos(dTheta) - fRadius;
}

namespace {
    class ICounterClockwisePredicate {
        COMET::IHoughLine* fLine;
    public:
        ICounterClockwisePredicate(COMET::IHoughLine* l) : fLine(l) {}
        bool operator () (const COMET::IHoughPoint* a, const COMET::IHoughPoint* b) {
            return fLine->GetLength(a) < fLine->GetLength(b);
        }
    };
}

void COMET::IHoughLine::SortPoints() {
    fSorted = true;
    fPoints.sort(ICounterClockwisePredicate(this));
}

COMET::IHoughLines::IHoughLines(double xmin, double ymin, 
                         double xmax, double ymax,
                         double lineWidth,
                         double overSampling,
                         double coherenceLength) {
    COMETInfo("HoughLines Space "
              << "[ " << xmin/unit::mm << " mm -- " 
              << xmax/unit::mm << " mm]"
              << " by [" << ymin/unit::mm << " mm -- "
              << ymax/unit::mm << " mm]");

    fCenterX = 0.5*(xmax+xmin);
    fCenterY = 0.5*(ymax+ymin);

    fLineWidth = std::abs(lineWidth);
    fLineSpacing = std::abs(fLineWidth/overSampling);
    coherenceLength = std::abs(coherenceLength);

    double cornerX = (xmax-fCenterX);
    double cornerY = (ymax-fCenterY);
    double radius = sqrt(cornerX*cornerX + cornerY*cornerY);
    fRadii = 2*int(radius/fLineSpacing)+3;
    fRadiusOffset = fRadii / 2;

    // Place a minimum limit on the coherence length.  If it is less than
    // this, then assume that the default value was intended.
    if (coherenceLength<4*fLineWidth) {
        coherenceLength = fLineWidth*(fRadii-2)/M_PI;
    }
    // Place a maximum limit on the coherence length since larger values don't
    // make since with the granularity of this transform.  Larger values just
    // waste memory and CPU.
    double maxLength = fLineWidth*radius/fLineSpacing;
    if (coherenceLength>maxLength) coherenceLength = maxLength;

    double approxAngleWidth = fLineWidth/coherenceLength;
    fAngles = int(2*(1+M_PI_2/approxAngleWidth) + 1);
    fAngleOffset = fAngles/2;
    fAngleSpacing = M_PI/(fAngles-1);

    COMETInfo("    Angular Binning: " << fAngles
              << " (" << fAngleSpacing/unit::degree << " degree)");

    COMETInfo("    Radial Binning:  " << fRadii
              << " (" << fLineWidth/unit::mm << " mm wide and"
              << " " << fLineSpacing/unit::mm << " mm apart)");

    // Create all of the lines.
    for (int a = 0; a<fAngles; ++a) {
        for (int r = 0; r<fRadii; ++r) {
            double angle = GetAngleValue(a);
            double radius = GetRadiusValue(r);
            fLines.push_back(COMET::IHoughLine(this,angle,radius));
        }
    }
}

COMET::IHoughLines::~IHoughLines() {
    for (Points::iterator p = fPoints.begin(); p != fPoints.end(); ++p) {
        delete (*p);
    }
}

int COMET::IHoughLines::DigitizeAngle(double a) {
    double bin = a/fAngleSpacing;
    return TMath::Nint(bin);
}

int COMET::IHoughLines::DigitizeRadius(double r) {
    double bin = r/fLineSpacing;
    return TMath::Nint(bin);
}

void COMET::IHoughLines::AddPoint(COMET::IHoughPoint* point) {
    point->fAngle = std::atan2(point->GetY()-fCenterY,point->GetX()-fCenterX);
    fPoints.push_back(point);
    
    for (int angleIndex=0; angleIndex<fAngles; ++angleIndex) {
        double angle = GetAngleValue(angleIndex);
        double radius = std::cos(angle)*(point->GetY()-fCenterY) 
            -std::sin(angle)*(point->GetX()-fCenterX);
        int angleBin = GetAngleBin(angleIndex);

        // Find the range of hough bins that will contain the point.  This
        // calculation is slightly obtuse since the bins are coarse wrt the
        // line width.  The line spacing is subtracted from the line width to
        // keep the oversampling fraction correct, but then a small tolerance
        // is added back to the width to force the digitized width to be
        // rounded up when the point falls on an exact bin center.
        double halfWidth = (fLineWidth-fLineSpacing)/2.0 + 1E-2*fLineSpacing;
        int highRadiusBin = DigitizeRadius(radius+halfWidth);
        int lowRadiusBin = DigitizeRadius(radius-halfWidth);

        for (int radiusBin=lowRadiusBin; 
             radiusBin<=highRadiusBin; 
             ++radiusBin) {
            int lineIndex = GetLinesIndex(angleBin,radiusBin);
            fLines[lineIndex].AddPoint(point);
        }
    }
}

void COMET::IHoughLines::Clear() {
    for (Points::iterator p = fPoints.begin(); p != fPoints.end(); ++p) {
        delete (*p);
    }
    fPoints.clear();

    for (Lines::iterator r = fLines.begin(); r != fLines.end(); ++r) {
        r->Clear();
    }
}

TH2F* COMET::IHoughLines::CreateHistogram(const char* name) {
    double maxAngle = M_PI_2;
    double maxRadius = (fAngles-fAngleOffset)*fLineSpacing;
    TH2F* hist = new TH2F(name,"Hough Transform Weights",
                          fAngles,-maxAngle,maxAngle,
                          fRadii,-maxRadius,maxRadius);
    
    for (int a = 0; a<fAngles; ++a) {
        for (int r =0; r<fRadii; ++r) {
            int lineIndex = GetLinesIndex(a-fAngleOffset,r-fRadiusOffset);
            double weight = fLines[lineIndex].GetWeight();
            hist->SetBinContent(a,r,weight);
       }
    }

    return hist;
}

namespace {
    bool LineWeightPredicate(COMET::IHoughLine* a, COMET::IHoughLine* b) {
        return a->GetWeight() < b->GetWeight();
    }
}

// For each point, find the most significant line that it is part of, and then
// remove it from every other line to which it contributes.
void COMET::IHoughLines::RemoveShared(double threshold) {
    std::vector<COMET::IHoughLine*> sortedLines;

    // Build a list of all lines that have more than three points in the line.
    for (COMET::IHoughLines::Lines::iterator ln = fLines.begin();
         ln != fLines.end();
         ++ln) {
        if (ln->GetPoints().size()<3) {
            // Reject really short lines.  
            ln->RemovePoints(ln->GetPoints().begin(),
                             ln->GetPoints().end());
            continue;
        }
        sortedLines.push_back(&(*ln));
    }

    // Sort the line by weight.  The least significant lines come first.
    std::sort(sortedLines.begin(), sortedLines.end(), LineWeightPredicate);

    // Pop off the most significant line remaining in the sorted lines vector.
    // Remove the points from all of the other lines they are a member of.  At
    // the end of this process, a point will belong to a single line, and that
    // line will be the most significant line that contained the point.
    while(sortedLines.size()>1) {
        COMET::IHoughLine* line = sortedLines.back();
        sortedLines.pop_back();
        if (line->GetPoints().size()<3) {
            line->RemovePoints(line->GetPoints().begin(),
                               line->GetPoints().end());
            continue;
        }
        for(COMET::IHoughLine::Points::iterator p = line->GetPoints().begin();
            p != line->GetPoints().end();
            ++p) {
            if (threshold < 0.999) (*p)->RemoveFromLines(threshold);
            else  (*p)->KeepLine(line);
        }
        std::sort(sortedLines.begin(), sortedLines.end(), LineWeightPredicate);
    }
}