// ============================================================================
//  quartic_hmc.cpp — single-file HMC code for the pure quartic matrix model
//                    (teaching version)
//
//      S(Phi) = N * [ (b/2) Tr(Phi^2)  +  c Tr(Phi^4) ]
//
//  Phi is a single N x N Hermitian matrix. This is the simplest nontrivial
//  Hermitian one-matrix model: for c = 0 it is exactly solvable (Gaussian
//  matrix model, Wigner semicircle eigenvalue distribution), and turning on
//  c drives the eigenvalue distribution from one lump ("one-cut") into two
//  ("two-cut", a double well) once b is negative enough — a cheap example of
//  a large-N phase transition.
//
//  Everything lives in this one file, organised top to bottom exactly in the
//  order the program uses it:
//
//      1. random numbers
//      2. parameters (defaults + "key=value" command line overrides)
//      3. physics: Action(), Force() = dS/dPhi, and a sanity check of the two
//      4. HMC update: initial condition, one leapfrog trajectory
//      5. observables: Tr(Phi^2) and the eigenvalues of Phi
//      6. output files
//      7. main(): thermalisation loop + production loop
//
//  Usage:
//      make
//      ./quartic_hmc                        # run with the defaults below
//      ./quartic_hmc N=30 c=2.0 ntraj=5000   # override any parameter
// ============================================================================

#include <armadillo>
#include <string>
#include <random>
#include <iostream>
#include <fstream>
#include <sstream>
#include <filesystem>
#include <cmath>
#include <stdexcept>

using namespace std;
using namespace arma;
namespace fs = std::filesystem;

// ============================================================================
//  1. RANDOM NUMBERS
// ============================================================================
static std::mt19937_64 gen;                       // the one RNG used everywhere
static std::uniform_real_distribution<> uni01(0.0, 1.0);

static void setup_rng(unsigned long seedval)
{
    if (seedval == 0) {
        std::random_device rd;          // fresh entropy from the OS
        gen.seed(rd());
    } else {
        gen.seed(seedval);              // reproducible run, useful for debugging
    }
}

static double normal_random(double mean, double sigma)
{
    std::normal_distribution<double> dist(mean, sigma);
    return dist(gen);
}

// random Hermitian size x size matrix with entries distributed so that
//     P(M) ~ exp( -Tr(M^2) / (2 sigma^2) )
// diagonal ~ Normal(0,sigma); off-diagonal real & imaginary parts each
// ~ Normal(0, sigma/sqrt(2)), so |M_ij|^2 has the same variance sigma^2.
static cx_mat randomHermitian(int size, double sigma)
{
    cx_mat M(size, size);
    for (int i = 0; i < size; i++) {
        M(i, i) = normal_random(0.0, sigma);
        for (int j = i + 1; j < size; j++) {
            double re = normal_random(0.0, sigma / std::sqrt(2.0));
            double im = normal_random(0.0, sigma / std::sqrt(2.0));
            M(i, j) = cx_double(re, im);
            M(j, i) = std::conj(M(i, j));
        }
    }
    return M;
}

// ============================================================================
//  2. PARAMETERS
//     defaults below; every one can be overridden on the command line as
//     key=value, e.g.  ./quartic_hmc N=30 b=-2.0 ntraj=5000
// ============================================================================
static int    N     = 50;      // matrix size
static double b     = -4.0;    // quadratic coupling
static double c     =  1.0;    // quartic coupling

static int    nsteps      = 4;      // leapfrog steps per HMC trajectory
static double stepLength  = 0.01;   // leapfrog step size epsilon

static int    ntherm      = 1000;   // thermalisation trajectories (not measured)
static int    ntraj       = 20000;  // production trajectories (measured)
static int    nmeas       = 1;      // measure S, Tr(Phi^2) every ... trajectories
static int    neig        = 10;     // measure the eigenvalue spectrum every ...
static int    print_every = 500;    // progress line every ... trajectories

static unsigned long seed = 0;             // 0 = pick a random seed
static std::string   outdir = "data";      // output directory

static void set_param(const std::string& key, const std::string& val)
{
    if      (key == "N")           N = std::stoi(val);
    else if (key == "b")           b = std::stod(val);
    else if (key == "c")           c = std::stod(val);
    else if (key == "nsteps")      nsteps = std::stoi(val);
    else if (key == "steplength")  stepLength = std::stod(val);
    else if (key == "ntherm")      ntherm = std::stoi(val);
    else if (key == "ntraj")       ntraj = std::stoi(val);
    else if (key == "nmeas")       nmeas = std::stoi(val);
    else if (key == "neig")        neig = std::stoi(val);
    else if (key == "print_every") print_every = std::stoi(val);
    else if (key == "seed")        seed = std::stoul(val);
    else if (key == "outdir")      outdir = val;
    else throw std::runtime_error("unknown parameter '" + key +
                                   "' (check the spelling against the top of this file)");
}

static void parse_command_line(int argc, char* argv[])
{
    for (int i = 1; i < argc; i++) {
        std::string arg = argv[i];
        size_t eq = arg.find('=');
        if (eq == std::string::npos)
            throw std::runtime_error("bad argument '" + arg +
                                      "' (expected key=value, e.g. N=30)");
        set_param(arg.substr(0, eq), arg.substr(eq + 1));
    }
}

// ============================================================================
//  DYNAMICAL VARIABLES
// ============================================================================
static cx_mat Phi;   // the N x N Hermitian matrix being simulated
static cx_mat Pi;    // its conjugate momentum (meaningful only mid-trajectory)

// ============================================================================
//  3. PHYSICS — Action(), Force() = dS/dPhi, and a sanity check of the two
//
//  HMC needs the gradient F = dS/dPhi, defined so that for any small
//  Hermitian "direction" E,
//
//      S(Phi + eps*E) - S(Phi - eps*E) = 2*eps*Tr(F*E) + O(eps^3) .
//
//  Using the standard matrix-calculus results
//
//      d Tr(Phi^2)/dPhi = 2 Phi ,        d Tr(Phi^4)/dPhi = 4 Phi^3
//
//  (easy to check by writing out indices and using cyclicity of the trace)
//  gives
//
//      Force = dS/dPhi = N * ( b*Phi + 4*c*Phi^3 ) .
//
//  check_force() verifies this formula against a numerical derivative of
//  Action() at startup — good practice for any HMC code: always test the
//  force before trusting a run.
// ============================================================================
static double Action()
{
    // Tr(Phi^2) and Tr(Phi^4) are real for Hermitian Phi; real(...) just
    // discards the rounding-level imaginary part.
    double TrPhi2 = real(trace(Phi * Phi));
    double TrPhi4 = real(trace(Phi * Phi * Phi * Phi));
    return N * (0.5 * b * TrPhi2 + c * TrPhi4);
}

static cx_mat Force()
{
    cx_mat Phi3 = Phi * Phi * Phi;
    cx_mat F = N * (b * Phi + 4.0 * c * Phi3);

    // F is already Hermitian in exact arithmetic; this just removes
    // floating-point rounding noise so Phi stays numerically Hermitian.
    return (F + F.t()) / 2.0;
}

static void check_force()
{
    cx_mat Phi_backup = Phi;                // do not disturb the real run

    Phi = randomHermitian(N, 0.7);          // random O(1) test configuration
    cx_mat Phi_test = Phi;
    cx_mat E = randomHermitian(N, 1.0);     // random Hermitian test direction

    double analytic = real(trace(Force() * E));

    const double eps = 1e-5;
    Phi = Phi_test + eps * E;  double S_plus  = Action();
    Phi = Phi_test - eps * E;  double S_minus = Action();
    double numeric = (S_plus - S_minus) / (2.0 * eps);

    Phi = Phi_backup;                       // restore the real configuration

    double err = std::abs(analytic - numeric) / std::max(1.0, std::abs(numeric));
    std::cout << "force check: analytic = " << analytic
              << ", numeric = " << numeric
              << ", relative error = " << err << std::endl;
    if (err > 1e-4)
        std::cout << "  WARNING: Force() and Action() disagree — "
                      "check the code before trusting this run!" << std::endl;
}

// ============================================================================
//  4. HMC UPDATE
//
//  HMC in a nutshell:
//    1. draw a random momentum Pi from the Gaussian heat bath exp(-Tr Pi^2/2)
//    2. evolve (Phi,Pi) for a short trajectory with a reversible,
//       volume-preserving integrator (leapfrog)
//    3. accept the proposed Phi with probability min(1, exp(-DeltaH)),
//       otherwise keep the old Phi
//  This samples Phi from exp(-S(Phi))/Z exactly, regardless of how large the
//  leapfrog step-size error is — the step size only affects the acceptance
//  rate, never the correctness of the result.
// ============================================================================
static void InitPhi()
{
    Phi = randomHermitian(N, 1.0 / std::sqrt(2.0));   // hot start
}

static void RefreshPi()
{
    Pi = randomHermitian(N, 1.0);           // heat bath: P(Pi) ~ exp(-Tr(Pi^2)/2)
}

static double Hamiltonian()
{
    return real(trace(Pi * Pi)) / 2.0 + Action();
}

// One HMC trajectory. Returns true if accepted; deltaH is the PROPOSED
// Delta H either way, so <exp(-deltaH)> -> 1 can be monitored honestly.
//
// Leapfrog integrates dPhi/dt = Pi, dPi/dt = -dS/dPhi with alternating
// half-kicks (momentum update) and drifts (position update); consecutive
// half-kicks between steps are merged into one full kick.
static bool leapfrogTrajectory(double& deltaH)
{
    RefreshPi();
    cx_mat Phi_old = Phi;
    double H_old = Hamiltonian();

    const double eps = stepLength;

    Pi -= 0.5 * eps * Force();                 // initial half-kick
    for (int step = 0; step < nsteps; step++) {
        Phi += eps * Pi;                        // drift
        if (step < nsteps - 1)
            Pi -= eps * Force();                // full kick between steps
        else
            Pi -= 0.5 * eps * Force();          // final half-kick
    }

    double H_new = Hamiltonian();
    deltaH = H_new - H_old;

    if (uni01(gen) < std::exp(-deltaH)) {
        return true;                            // accept: keep the new Phi
    } else {
        Phi = Phi_old;                          // reject: restore the old Phi
        return false;
    }
}

// ============================================================================
//  5. OBSERVABLES — exactly what gets measured and printed: the eigenvalues
//     of Phi, Tr(Phi^2), and (for free, and as the standard first sanity
//     check of any HMC run) the action S itself.
// ============================================================================
static double measure_TrPhi2()
{
    return real(trace(Phi * Phi)) / N;
}

static vec measure_eigenvalues()
{
    return eig_sym(Phi);        // ascending eigenvalues (arma default)
}

// ============================================================================
//  6. OUTPUT FILES
//
//      obs_<run_name>.dat   columns: trajectory   S   Tr(Phi^2)/N
//      eig_<run_name>.dat   columns: trajectory   lambda_1 ... lambda_N
//      info_<run_name>.txt  human-readable summary of the run
//
//  No header line in the .dat files, so they load directly with e.g.
//  numpy.loadtxt("obs_....dat").
// ============================================================================
static std::string run_name()
{
    std::ostringstream os;
    os << "N" << N << "_b" << b << "_c" << c;
    return os.str();
}

static std::string obs_path()  { return outdir + "/obs_"  + run_name() + ".dat"; }
static std::string eig_path()  { return outdir + "/eig_"  + run_name() + ".dat"; }
static std::string info_path() { return outdir + "/info_" + run_name() + ".txt"; }

static void open_output_files()
{
    fs::create_directories(outdir);
    std::ofstream(obs_path(), std::ios::out).close();   // truncate / create
    std::ofstream(eig_path(), std::ios::out).close();
    std::cout << "writing to " << outdir << "/  (files *_" << run_name() << ")"
              << std::endl;
}

static void append_observables(int traj, double S, double TrPhi2)
{
    std::ofstream out(obs_path(), std::ios::app);
    out.precision(10);
    out << traj << " " << S << " " << TrPhi2 << "\n";
}

static void append_eigenvalues(int traj, const vec& eig)
{
    std::ofstream out(eig_path(), std::ios::app);
    out.precision(8);
    out << traj;
    for (int k = 0; k < (int)eig.n_elem; k++) out << " " << eig(k);
    out << "\n";
}

static void write_info_file()
{
    std::ofstream out(info_path());
    out << "quartic Hermitian matrix model - HMC run\n";
    out << "=========================================\n\n";
    out << "action:  S(Phi) = N * [ (b/2) Tr(Phi^2) + c Tr(Phi^4) ]\n\n";
    out << "N            = " << N << "\n";
    out << "b            = " << b << "\n";
    out << "c            = " << c << "\n";
    out << "nsteps       = " << nsteps << "\n";
    out << "stepLength   = " << stepLength << "\n";
    out << "ntherm       = " << ntherm << "\n";
    out << "ntraj        = " << ntraj << "\n";
    out << "nmeas        = " << nmeas << "\n";
    out << "neig         = " << neig << "\n";
    out << "seed         = " << seed << "\n\n";
    out << "files:\n";
    out << "  obs_" << run_name() << ".dat : trajectory  S  Tr(Phi^2)/N\n";
    out << "  eig_" << run_name() << ".dat : trajectory  lambda_1 ... lambda_N"
                                     " (ascending)\n";
    out.close();
    std::cout << "run info written to " << info_path() << std::endl;
}

// ============================================================================
//  7. main(): thermalisation + production
// ============================================================================
int main(int argc, char* argv[])
{
  try {
    parse_command_line(argc, argv);
    setup_rng(seed);

    std::cout << "quartic matrix model:  S(Phi) = N*[ (b/2) Tr(Phi^2) + c Tr(Phi^4) ]\n"
              << "N=" << N << "  b=" << b << "  c=" << c
              << "  nsteps=" << nsteps << "  stepLength=" << stepLength
              << std::endl;

    InitPhi();
    check_force();                 // one-time sanity check of Force() vs Action()

    open_output_files();
    write_info_file();

    // ------------------------------------------------------------------------
    //  THERMALISATION: run the Markov chain without recording data, so Phi
    //  forgets its (arbitrary) starting point and settles onto the
    //  equilibrium distribution exp(-S(Phi)).
    // ------------------------------------------------------------------------
    std::cout << "thermalising (" << ntherm << " trajectories)..." << std::endl;
    double accepted = 0;
    for (int i = 1; i <= ntherm; i++) {
        double deltaH;
        if (leapfrogTrajectory(deltaH)) accepted++;

        if (print_every > 0 && i % print_every == 0)
            std::cout << "  therm " << i << "/" << ntherm
                      << "   acceptance=" << accepted / i << std::endl;
    }

    // ------------------------------------------------------------------------
    //  PRODUCTION: measure observables and write them to disk
    // ------------------------------------------------------------------------
    std::cout << "production (" << ntraj << " trajectories)..." << std::endl;
    accepted = 0;
    double sum_expdH = 0;

    for (int i = 1; i <= ntraj; i++) {
        double deltaH;
        if (leapfrogTrajectory(deltaH)) accepted++;
        sum_expdH += std::exp(-deltaH);

        if (i % nmeas == 0)
            append_observables(i, Action(), measure_TrPhi2());
        if (i % neig == 0)
            append_eigenvalues(i, measure_eigenvalues());

        if (print_every > 0 && i % print_every == 0)
            std::cout << "  i=" << i << "/" << ntraj
                      << "   acceptance=" << accepted / i
                      << "   <exp(-deltaH)>=" << sum_expdH / i
                      << "   (should -> 1)" << std::endl;
    }

    std::cout << "done.  final acceptance = " << accepted / ntraj
              << "   <exp(-deltaH)> = " << sum_expdH / ntraj
              << "   (should be close to 1 — the standard HMC correctness check)"
              << std::endl;
    return 0;
  }
  catch (const std::exception& e) {
      std::cerr << "ERROR: " << e.what() << std::endl;
      return 1;
  }
}