#include <RAT/DropoutFitter.hh>
#include <RAT/Optimiser.hh>
#include <RAT/SimulatedAnnealing.hh>
#include <TMath.h>

namespace RAT {
    DropoutFitter::DropoutFitter(){
        fSampleCount = 0;
        fOptimiser = new Optimisers::SimulatedAnnealing();
        fOptimiser->SetComponent(this);
        fOptimiser->Initialise("");
        // Set up inital params for fit
        double initial_vals[] = {3, 4000, 42, 15};
        double initial_err[] = {1.5, 150, 20, 10};
        fParams.resize(4);
        fErrors.resize(4);
        for(int i=0; i < 4; i++) {
            fParams[i] = initial_vals[i];
            fErrors[i] = initial_err[i];
        }
    }

    DropoutFitter::~DropoutFitter() {
        delete fOptimiser;
    }

    vector<double> DropoutFitter::GetParams() const {
        return fParams;
    }

    void DropoutFitter::SetParams(const vector<double>& params) {
        fParams = params;
    }

    void DropoutFitter::SetParam(const int& iparam, const double& value) {
        fParams[iparam] = value;
    }

    vector<double> DropoutFitter::GetPositiveErrors() const {
        return fErrors;
    }

    vector<double> DropoutFitter::GetNegativeErrors() const {
        return fErrors;
    }

    void DropoutFitter::SetError(const int& iparam, const double& value) {
        fErrors[iparam] = value;
    }

    void DropoutFitter::SetFOM(const std::string& fomName, const double fom) {
        fFoms[fomName] = fom;
    }

    void DropoutFitter::SetTemperature(const int& temp) {
        fTemperature = temp;
        fOptimiser->SetD("startingTemperature", fTemperature);

    }

    void DropoutFitter::SetIterations(const int& iterations) {
        fIterations = iterations;
        fOptimiser->SetI("iterations", fIterations);
    }

    void DropoutFitter::PerformFit() {
        fSampleCount = 0;
        while(fSampleCount < fIterations) {
            fOptimiser->Maximise();
        }
        SetParams(fOptimiser->GetParams());
    }

    double DropoutFitter::Chi2(const vector<double>& y, const vector<double>& yp) {
        size_t len = y.size();
        if (yp.size() < len) {
            len = yp.size();
        }
        if(len == 0) {
            return -1;
        }
        double chi2 = 0;
        for(size_t i=0; i < len; i++) {
            // Don't let the prediction be less than 1e-6
            // Otherwise it can get really small (e.g. 1e-324) and cause issues
            double prediction = yp[i] > 1e-6 ? yp[i] : 1e-6;
            double diff = y[i] - prediction;
            chi2 += diff*diff/prediction;
        }
        return fNormalization*chi2/(double)len;
    }

    double DropoutFitter::LogLikelihood(const vector<double>& y, const vector<double>& yp) {
        size_t len = y.size();
        if (yp.size() < len) {
            len = yp.size();
        }
        double ll = 0;
        for(size_t i = 0; i < len; i++) {
            double pred = yp[i];
            double obs = y[i];

            // Very small numbers can produce nonsesical log values.
            // So it's best to just ignore them
            if( pred < 1e-7 && obs < 1e-7) {
                continue;
            }
            ll += log(pow(pred, obs));
        }
        return ll;
    }

    double DropoutFitter::operator()(const vector<double>& params) {
        // Evaluates the (log)likelihood of the given parameters

        double ll = 0;
        double rate = params.at(0);
        double pos = params.at(1);
        double sep = params.at(2);
        double sigma = params.at(3);

        // Calculate priors first
        double p = Prior(rate, pos, sep, sigma);
        if(p == 0) {
            return -1e6;
        }
        ll += log(p);

        // If the prior is non-zero then carry on with the fit
        fSampleCount+=1;
        // Calculate the predicted dropout "spectrum"
        vector<double> yp = DropoutModel(xax, rate, pos, sep, sigma);

        // Evaluate the likelihood of that spectrum
        ll += LogLikelihood(yvals, yp);
        return ll;
    }

    // TODO consider adding return y values to input args as optimization
    vector<double> DropoutFitter::DropoutModel(const vector<double> &x,
                                               const double& rate,
                                               const double& pos,
                                               const double& separation,
                                               const double& sigma) {
        // Given a few parameters this function creates a predicted
        // dropout "spectrum".
        // The basics of the model is that the dropout one expects for any
        // random event is poissonian around some rate. And the CAEN baseline
        // that is observed will be a guassian (with some position and sigma).
        // So you expect several gaussians each seperated by a constant amount
        // With the relative sizes given by the possiion rate.
        const size_t N_PEAKS = 10;
        vector<double> ret;
        ret.resize(x.size());
        double peak_locations[N_PEAKS];
        double poisson_scale[N_PEAKS];
        double clip_scaling[N_PEAKS];
        double scales[N_PEAKS];
        double sum =0;


        for(size_t i =0; i < N_PEAKS; i++) {
            peak_locations[i] = pos - separation*i;
            poisson_scale[i] = TMath::Poisson(i, rate);
            clip_scaling[i] = ClipProbability(peak_locations[i], sigma);
        }

        for(size_t i =0; i < N_PEAKS; i++) {
            scales[i] = poisson_scale[i]*clip_scaling[i];
            sum += scales[i];
        }
        // Normalize
        for(size_t i =0; i < N_PEAKS; i++) {
            scales[i] /= sum;
        }

        for(size_t i = 0; i < x.size(); i++) {
            ret[i] = 0;
            for(size_t j=0; j <N_PEAKS; j++) {
                ret[i] += scales[j]*TMath::Gaus(x[i], peak_locations[j], sigma, true);
            }
        }
        return ret;
    }

    double DropoutFitter::Prior(const double& rate, const double& pos,
                                const double& sep, const double& sigma) {
        // Evaluates if the given parameters are reasonable for the SNO+ detector
        if (rate < 0 || rate > 10) {
            return 0;
        }
        if(pos < 0 or pos > 5000) {
            return 0;
        }
        if(sigma <= 0 || sigma > 40) {
            return 0;
        }
        if(sep <= 0 || sep < 2*sigma || sep > 75) {
            return 0;
        }
        return 1;
    }

    double DropoutFitter::ClipProbability(const double& mu, const double& sigma) {
        // Returns fraction of a normalized gaussian that will get clipped by the
        // CAEN's upper limit of 4096

        static double inv_sqrt_two = 1./sqrt(2); // Minor optimization
        double distance = (4096 - mu)/sigma;

        // For ROOT's error function implemenation the sigma is sqrt(1/2) so I have
        // to scale my effective distance appropriately
        distance = distance*inv_sqrt_two;
        if (distance < 0) {
            return 1 - (TMath::Erf(-1*distance)/2.0 + 0.5);
        }
        return 0.5 + TMath::Erf(distance)/2.0;
    }

}