// ch13_linear_fit.cpp
// Companion code for Chapter 13: Optimization and Data Fitting.
// Standalone C++11; standard library only; console output without plotting.
// Compile: c++ -std=c++11 -O2 -Wall -Wextra -pedantic ch13_linear_fit.cpp -o ch13_linear_fit
// Edit the example settings in main(); routines above main() can be reused.
// Fixed seeds reproduce runs with the same standard-library implementation.
// Random distributions need not give identical sequences across compilers/languages.

#include <algorithm>
#include <array>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <numeric>
#include <random>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

using Vec = std::vector<double>;
using Mat = std::vector<Vec>;
using Function = std::function<double(const Vec&)>;
using ScalarFunction = std::function<double(double)>;
const double pi = std::acos(-1.0);

void require(bool condition, const std::string& message) {
    if (!condition) throw std::invalid_argument(message);
}
void finite_vector(const Vec& x) {
    require(!x.empty(), "A nonempty vector is required.");
    for (double v : x) require(std::isfinite(v), "Data must be finite.");
}
void check_bounds(const Vec& lo, const Vec& hi) {
    finite_vector(lo); finite_vector(hi);
    require(lo.size() == hi.size(), "Bounds must have matching dimensions.");
    for (std::size_t j=0; j<lo.size(); ++j)
        require(hi[j]>lo[j] && std::isfinite(hi[j]-lo[j]), "Require finite upper > lower.");
}
double objective_value(double value) {
    require(std::isfinite(value), "Objective must return a finite real scalar.");
    return value;
}
double clip(double x, double lo, double hi) { return std::max(lo,std::min(hi,x)); }
Vec linspace(double a, double b, int n) {
    require(n>=2, "Grid needs at least two points.");
    Vec x(n);
    for (int i=0; i<n; ++i) x[i]=a+(b-a)*i/(n-1);
    return x;
}
struct Random {
    std::mt19937 engine;
    std::uniform_real_distribution<double> uniform{0.0,1.0};
    std::normal_distribution<double> normal{0.0,1.0};
    explicit Random(unsigned int seed) : engine(seed) {}
    double u() { return uniform(engine); }
    double z() { return normal(engine); }
    int index(int n) { return std::uniform_int_distribution<int>(0,n-1)(engine); }
};
struct Result {
    Vec point;
    double value=0.0;
    long long evaluations=0;
    int iterations=0;
    std::string status;
};
void print_result(const Result& r) {
    std::cout << std::setprecision(12);
    for (std::size_t j=0; j<r.point.size(); ++j)
        std::cout << "p(" << j+1 << ") = " << r.point[j] << '\n';
    std::cout << "Objective value: " << r.value << "\nIterations/updates: " << r.iterations
              << "\nFunction calls: " << r.evaluations << "\nStatus: " << r.status << '\n';
}

struct LineFit {
    double intercept=0, slope=0, sse=0, residual_std=0, chi2=0;
    int dof=0; bool known_sigma=false;
    Mat covariance;
    Vec standard_errors, residuals;
};
// Empty sigma means an unknown common noise variance; otherwise sigma contains
// known absolute standard deviations. Assumes independent y errors, negligible x errors.
LineFit linear_fit(const Vec& x, const Vec& y, const Vec& sigma=Vec()) {
    finite_vector(x); finite_vector(y); std::size_t n=x.size();
    require(n>=3 && y.size()==n, "Require matching x,y and at least three observations.");
    LineFit r; r.known_sigma=!sigma.empty(); Vec w(n,1.0);
    if (r.known_sigma) {
        finite_vector(sigma); require(sigma.size()==n, "One uncertainty per observation is required.");
        for (std::size_t i=0;i<n;++i) {
            require(sigma[i]>0, "Uncertainties must be positive."); w[i]=1/(sigma[i]*sigma[i]);
            require(std::isfinite(w[i]) && w[i]>0, "Numerically unusable weights.");
        }
    }
    double W=0,xb=0,yb=0;
    for (std::size_t i=0;i<n;++i) { W+=w[i]; xb+=w[i]*x[i]; yb+=w[i]*y[i]; }
    require(std::isfinite(W) && W>0, "Weight sum overflow."); xb/=W; yb/=W;
    double Sxx=0,Sxy=0;
    for (std::size_t i=0;i<n;++i) { double dx=x[i]-xb; Sxx+=w[i]*dx*dx; Sxy+=w[i]*dx*(y[i]-yb); }
    require(std::isfinite(Sxx) && Sxx>0 && std::isfinite(Sxy+xb+yb), "No usable x spread or data scaling.");
    r.slope=Sxy/Sxx; r.intercept=yb-r.slope*xb; r.dof=static_cast<int>(n)-2;
    r.residuals.resize(n);
    for (std::size_t i=0;i<n;++i) {
        double e=y[i]-(yb+r.slope*(x[i]-xb)); r.residuals[i]=e; r.sse+=e*e;
        if (r.known_sigma) r.chi2+=(e/sigma[i])*(e/sigma[i]);
    }
    r.covariance={{1/W+xb*xb/Sxx,-xb/Sxx},{-xb/Sxx,1/Sxx}};
    if (!r.known_sigma) {
        r.chi2=std::numeric_limits<double>::quiet_NaN();
        for (Vec& row:r.covariance) for (double& a:row) a*=r.sse/r.dof;
    }
    r.standard_errors={std::sqrt(r.covariance[0][0]),std::sqrt(r.covariance[1][1])};
    r.residual_std=std::sqrt(r.sse/r.dof); return r;
}

int main() {
  try {
    const unsigned int seed = 1;
    Random rng(seed);
    // Alternative: comment out the previous line and seed from system entropy.
    // Random rng(std::random_device{}());
    Vec x=linspace(0,10,15), y(x.size()),sigma_y(x.size(),.5);
    for (std::size_t i=0;i<x.size();++i) y[i]=1+.5*x[i]+.5*rng.z();
    // Replace x,y,sigma_y above with measured data. Empty sigma_y requests OLS.
    LineFit r=linear_fit(x,y,sigma_y);
    std::cout << std::setprecision(12) << "Fit y = a0 + a1*x\na0 = " << r.intercept << " +/- " << r.standard_errors[0]
              << "\na1 = " << r.slope << " +/- " << r.standard_errors[1]
              << "\nResidual sum of squares: " << r.sse << "\nResidual standard deviation: " << r.residual_std
              << "\nDegrees of freedom: " << r.dof << "\nChi-squared: " << r.chi2 << '\n';
    std::cout << (r.known_sigma ? "Covariance uses supplied absolute uncertainties.\n" : "Covariance scaled by residual variance.\n")
              << "Parameter errors are standard errors, not confidence intervals.\n";
    return 0;
  } catch (const std::exception& e) {
    std::cerr << "Error: " << e.what() << '\n'; return 1;
  }
}