/*
  Copyright (C) 2004-2007,2009 Benjamin Redelings

  This file is part of BAli-Phy.

  BAli-Phy is free software; you can redistribute it and/or modify it under
  the terms of the GNU General Public License as published by the Free
  Software Foundation; either version 2, or (at your option) any later
  version.

  BAli-Phy is distributed in the hope that it will be useful, but WITHOUT ANY
  WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  along with BAli-Phy; see the file COPYING.  If not see
  <http://www.gnu.org/licenses/>.  */

#include <cmath>
#include "util/assert.hh"
#include "statistics.H"

using std::valarray;
using std::vector;
using std::pair;

//FIXME - add 'hat'?
namespace statistics {

    using std::abs;

    double fraction(const valarray<bool>& v) {
        return fraction(count(v),v.size(),0);
    }
  
    double Pr(const valarray<bool>& v) {
        return fraction(v);
    }
  
    double odds(const valarray<bool>& v) {
        double y = count(v);
        double n = v.size() - count(v);
        return y/n;
    }

    double log_odds(const valarray<bool>& v) {
        return log(odds(v));
    }

    double quantile_sorted(const vector<double>& values, double Q)
    {
        assert(0 <= Q and Q <= 1.0);
        assert(values.size() > 0);

        if (values.size() == 1)
            return values[0];

        // compute indices
        double index = Q*values.size() - 0.5;
        int index1 = (int)index;
        int index2 = index1+1;
        double p1 = index2-index;
        double p2 = 1.0-p1;

        // compute quantile
        if (index1 < 0)                      // index <= 0
            return values[0];
        else if (index2 >= values.size())   // index >= values.size() - 1
            return values.back();
        else {
            assert(index1 >= 0 and index1 < values.size());
            assert(index2 >= 0 and index2 < values.size());
            return p1*values[index1] + p2*values[index2];
        }
    }

    double quantile(vector<double> values, double Q)
    {
        assert(0 <= Q and Q <= 1.0);
        assert(values.size() > 0);

        if (values.size() == 1)
            return values[0];

        // sort values
        sort(values.begin(),values.end());

        return quantile_sorted(values,Q);
    }

    double quantile(const valarray<double>& values, double Q)
    {
        assert(0 <= Q and Q <= 1.0);
        assert(values.size() > 0);

        if (values.size() == 1)
            return values[0];

        // sort values
        vector<double> values2(values.size());
        for(int i=0;i<values.size();i++)
            values2[i] = values[i];

        return quantile(values2,Q);
    }

    double fraction_in_interval(vector<double> values, double L, double R)
    {
        assert(L <= R);
        assert(values.size() > 0);

        int count=0;
        for(int i=0;i<values.size();i++)
            if (L <= values[i] and values[i] <= R)
                count++;

        return double(count)/values.size();
    }

    double median(const valarray<double>& values)
    {
        return quantile(values,0.5);
    }

    double median(const vector<double>& values)
    {
        return quantile(values,0.5);
    }

    /// FIXME - this could be faster by a factor of 2 - we are sorting TWICE!
    pair<double,double> central_confidence_interval(const valarray<double>& values,double P) 
    {
        vector<double> values2(values.size());
        for(int i=0;i<values.size();i++)
            values2[i] = values[i];

        return central_confidence_interval(values2,P);
    }

    pair<double,double> central_confidence_interval(vector<double> values,double P) 
    {
        double alpha = (1.0-P)/2.0;

        sort(values.begin(),values.end());

        pair<double,double> interval;
        interval.first = quantile_sorted(values,alpha);
        interval.second = quantile_sorted(values,1.0-alpha);
    
        return interval;
    }

    /// FIXME - this could be faster by a factor of 2 - we are sorting TWICE!
    pair<double,double> HPD_confidence_interval(const valarray<double>& values,double P) 
    {
        vector<double> values2(values.size());
        for(int i=0;i<values.size();i++)
            values2[i] = values[i];

        return HPD_confidence_interval(values2,P);
    }


    int find_smallest_interval(const vector<double>& values, int delta)
    {
        const int N = values.size();
    
        int ways = N - delta - 1;

        // scan for the interval with the smallest width
        double width = values.back() - values[0];
        int best = 0;
        for(int i=0;i<ways;i++)
        {
            assert(i+delta < N);
            double w = values[i+delta] - values[i];
            if (w < width)
            {
                best = i;
                width = w;
            }
        }

        return best;
    }

    pair<double,double> HPD_confidence_interval(vector<double> values,double P) 
    {
        if (values.empty()) return pair<double,double>{0,-1};

        sort(values.begin(),values.end());

        int N = values.size();
        int delta = ceil( P*N )-1;

        int best = find_smallest_interval(values, delta);
    
        return pair<double,double>{values[best],values[best+delta]};
    }

    pair<double,double> mode(vector<double> values)
    {
        if (values.size() < 20)
            return pair<double,double>(0, -1);

        std::sort(values.begin(),values.end());

        const int delta = sqrt(values.size()*1.6);

        int best = find_smallest_interval(values, delta);
    
        double x1 = values[best];
        double x2 = values[best+delta];
    
        return pair<double,double>(0.5*(x1+x2),abs(x2-x1));
    }
  
    valarray<bool> add_pseudocount(const valarray<bool>& sample1,int pseudocount) {
        valarray<bool> sample2(sample1.size() + 2*pseudocount);
    
        int i=0;
        for(int j=0;j<pseudocount;i++,j++)
            sample2[i] = true;
    
        for(int j=0;j<sample1.size();i++,j++)
            sample2[i] = sample1[j];
    
        for(int j=0;j<pseudocount;i++,j++)
            sample2[i] = false;
    
        return sample2;
    }

    vector<int> total_times(const valarray<bool>& v) 
    {
        vector<int> total(v.size()+1);

        total[0] = 0;
        for(int i=0;i<v.size();i++) {
            total[i+1] = total[i];
            if (v[i]) total[i+1]++;
        }
      
        return total;
    }

    vector<int> regeneration_times(const valarray<bool>& v) 
    {
        vector<int> times;

        times.push_back(0);
        for(int i=1;i<v.size();i++) {
            if (v[i] and not v[i-1])
                times.push_back(i);
        }
        times.push_back(v.size());
      
        return times;
    }

    // Would it be faster to simply compute \sum x[i]*x[i+k]/(N-k) - mu^2?>
    // Well, the problem with this is that you never stop when rho[k] gets negative.

    vector<double> autocovariance(const vector<double>& x,unsigned max)
    {
        const int N = x.size();

        // specify sample size if unspecified
        if (max==0) 
            max = std::max(1+N/4,N-15);

        if (max >= N) 
            max = N-1;

        // compute mean of X
        double mean = 0;
        for(int i=0;i<N;i++)
            mean += x[i];
        mean /= N;

        // allocate covariances
        vector<double> rho(max);

        // Run iteration 0 separately - to use rho[0] in the limit calculation
        {
            int k=0;
            double total = 0;
            for(int i=0;i<N-k;i++)
                total += (x[i]-mean)*(x[i+k]-mean);

            rho[k] = total/(N-k);
        }

        // compute each autocorrelation rho[k]
        double limit = rho[0]*(0.01/N);
        for(int k=1;k<max;k++) 
        {
            double total = 0;
            for(int i=0;i<N-k;i++)
                total += (x[i]-mean)*(x[i+k]-mean);

            rho[k] = total/(N-k);

            if (rho[k] < 0 or (k > 3 and rho[k] < limit))
            {
                rho.resize(k);
                break;
            }
        }

        assert(not rho.empty());
        return rho;
    }

    vector<double> autocorrelation(const vector<double>& x,unsigned N)
    {
        vector<double> rho = autocovariance(x,N);
        const double V = rho[0];
        for(int i=0;i<rho.size();i++)
            rho[i] /= V;
        return rho;
    }

    double autocorrelation_time(const vector<double>& x,unsigned max)
    {
        if (x.size() < 2) return 1.0;

        vector<double> cv = autocovariance(x,max);

        double V = cv[0];

        double sum = 0;
        for(int i=1;i<cv.size();i++)
            sum += cv[i];

        return (1.0 + 2.0*sum/V);
    }

    double probability_x_less_than_y(const std::valarray<double>& x, const std::valarray<double>& y)
    {
        vector<double> x2(x.size());
        for(int i=0;i<x2.size();i++)
            x2[i] = x[i];
        sort(x2.begin(),x2.end());

        vector<double> y2(y.size());
        for(int i=0;i<y2.size();i++)
            y2[i] = y[i];
        sort(y2.begin(),y2.end());

        valarray<double> FX(y.size());
        int dx=0;
        for(int dy=0;dy<FX.size();dy++) {
            while((dx<x2.size()) and (x2[dx]<y2[dy]))
                dx++;
            FX[dy] = double(dx)/x2.size();
        }
        return FX.sum()/FX.size();
    }

    double bias(const vector<double>& x, double truth)
    {
        return average(x) - truth;
    }
  
    double overestimate(const vector<double>& xs, double truth)
    {
        int count = 0;
        for(double x:xs)
            if (x > truth)
                count++;

        double O = double(count)/xs.size();
        return O - 0.5;
    }
  
    double abs_error(const vector<double>& xs, double truth)
    {
        double E = 0;
        for(double x:xs)
            E += abs(x - truth);

        E /= xs.size();

        return E;
    }
  
    double rms_error(const vector<double>& xs, double truth)
    {
        double E = 0;
        for(double x:xs)
        {
            double e = x - truth;
            E += e*e;
        }

        E /= xs.size();
        return sqrt(E);
    }
}

