main.cpp

// =============================================================================
// main.cpp  —  Orquestador: PI-NSGA-II (Solo nuestro Algoritmo)
//   Modos:
//     ./build/pi_nsga2              → 1 corrida (comportamiento estándar)
//     ./build/pi_nsga2 --runs N     → N corridas independientes (análisis estadístico)
//
//   Genera CSVs de frentes de Pareto para análisis comparativo.
//   Ecuaciones: Laplace, Poisson, Helmholtz, Liouville, Sine-Gordon.
// =============================================================================

#include "pi_solver.hpp"
#include "numerical_solver.hpp"
#include <iostream>
#include <fstream>
#include <iomanip>
#include <string>
#include <vector>
#include <chrono>
#include <filesystem>
#include <algorithm>
#include <numeric>
#include <cstring>
#include <omp.h>

#include "common.hpp"
#include "pde_problems.hpp"
#include "pi_solver.hpp"

namespace fs = std::filesystem;

// ─── Guardar frente de Pareto como CSV ───────────────────────────────────────
template<typename Ind>
void save_pareto_csv(const std::vector<Ind>& pop,
                     const std::string& path,
                     const std::string& method,
                     const std::string& pde_label,
                     int dim)
{
    std::ofstream f(path);
    f << std::fixed << std::setprecision(10);
    f << "method,pde,dim,mse_domain,mse_boundary,tree_size,rank\n";
    for (auto& ind : pop) {
        if (ind.rank == 1)
            f << method << "," << pde_label << "," << dim << ","
              << ind.mse_domain << "," << ind.mse_boundary << ","
              << ind.tree_size << "," << ind.rank << "\n";
    }
}

// ─── Guardar Historial de Convergencia ────────────────────────────────────────
void save_convergence_csv(const std::vector<ConvergenceStats>& history,
                           const std::string& path)
{
    std::ofstream f(path);
    f << "gen,best_mse_domain,best_mse_boundary,best_total_mse\n";
    for (auto& s : history) {
        f << s.gen << "," << s.best_mse_domain << ","
          << s.best_mse_boundary << "," << s.best_total_mse << "\n";
    }
}

// ─── Guardar Mejor Expresión en LaTeX ────────────────────────────────────────
template<typename Ind>
void save_best_expression(const std::vector<Ind>& pop, const std::string& path) {
    const Ind* best = nullptr;
    double min_err = 1e18;
    for (auto& ind : pop) {
        if (ind.rank == 1) {
            double err = ind.mse_domain + ind.mse_boundary;
            if (err < min_err) { min_err = err; best = &ind; }
        }
    }
    if (best) {
        std::ofstream f(path);
        f << "$$ \\hat{u}(x,y) = ";
        auto simple_tree = best->tree->simplify(); // Limpieza antes de exportar
        simple_tree->print_latex(f);
        f << " $$" << std::endl;
    }
}

// ─── Guardar Rejilla de Evaluación (para plots 3D/1D) ─────────────────────────
template<typename Ind>
void save_best_grid(const std::vector<Ind>& pop, const PDEProblem& prob, const std::string& path) {
    const Ind* best = nullptr;
    double min_err = 1e18;
    for (auto& ind : pop) {
        if (ind.rank == 1) {
            double err = ind.mse_domain + ind.mse_boundary;
            if (err < min_err) { min_err = err; best = &ind; }
        }
    }
    if (best) {
        std::ofstream f(path);
        if (prob.dim == 1) {
            f << "x,u_exact,u_approx\n";
            for (int i = 0; i <= 100; ++i) {
                double x = (double)i / 100.0;
                f << x << "," << std::real(prob.exact(x, 0)) << "," << std::real(best->tree->eval(x, 0)) << "\n";
            }
        } else {
            f << "x,y,u_exact,u_approx\n";
            int N = 50;
            for (int i = 0; i <= N; ++i) {
                for (int j = 0; j <= N; ++j) {
                    double x = (double)i / N;
                    double y = (double)j / N;
                    f << x << "," << y << "," << std::real(prob.exact(x, y)) << "," << std::real(best->tree->eval(x, y)) << "\n";
                }
            }
        }
    }
}

// ─── Estadísticas de una corrida ─────────────────────────────────────────────
struct Stats {
    std::string method, pde;
    int    front_size  = 0;
    double best_domain = 1e18;
    double best_bnd    = 1e18;
    double mean_domain = 0.0;
    double mean_bnd    = 0.0;
    double runtime_s   = 0.0;
    double hypervolume = 0.0;
};

// ─── Hipervolumen 3D: mse_domain × mse_boundary × tree_size (normalizado) ─────
// Algoritmo: slice-by-slice sweep sobre el tercer objetivo (tree_size).
// Complejidad: O(n² log n), exacto para 3 objetivos.
// Referencia: Emmerich et al. (2006), WFG algorithm.
//
// Objetivos (todos a minimizar):
//   f1 = mse_domain     (ref: ref_dom)
//   f2 = mse_boundary   (ref: ref_bnd)
//   f3 = tree_size_norm (ref: 1.0, normalizado al máximo permitido)
// ─────────────────────────────────────────────────────────────────────────────

// Calcula HV 2D del conjunto de puntos 2D respecto al punto de referencia (rx, ry).
// Los puntos deben estar en el espacio (minimización); solo se consideran los que
// dominan al punto de referencia (xi < rx AND yi < ry).
static double hv2d(std::vector<std::pair<double,double>> pts, double rx, double ry)
{
    // Filtrar puntos que no dominan la referencia
    pts.erase(std::remove_if(pts.begin(), pts.end(),
        [rx, ry](const std::pair<double,double>& p){
            return p.first >= rx || p.second >= ry;
        }), pts.end());

    if (pts.empty()) return 0.0;

    // Ordenar por primer objetivo (ascendente)
    std::sort(pts.begin(), pts.end(),
              [](const std::pair<double,double>& a, const std::pair<double,double>& b){
                  return a.first < b.first;
              });

    // Sweepline: acumular área
    double hv = 0.0;
    double prev_x = pts[0].first;
    double cur_min_y = pts[0].second;

    for (size_t i = 1; i < pts.size(); ++i) {
        double x = pts[i].first;
        if (cur_min_y < ry)                         // solo si contribuye
            hv += (x - prev_x) * (ry - cur_min_y);
        cur_min_y = std::min(cur_min_y, pts[i].second);
        prev_x = x;
    }
    if (cur_min_y < ry)
        hv += (rx - prev_x) * (ry - cur_min_y);

    return hv;
}

template<typename Ind>
double compute_hypervolume(const std::vector<Ind>& pop,
                           double ref_dom  = 1e4,
                           double ref_bnd  = 1e4,
                           double ref_size = 1.0)   // tree_size normalizado
{
    // ── Recopilar puntos del frente de Pareto (rank == 1) ──────────────────
    struct Pt3 { double f1, f2, f3; };
    std::vector<Pt3> pts;

    // Determinar max tree_size para normalización
    double max_ts = 1.0;
    for (auto& ind : pop)
        if (ind.rank == 1)
            max_ts = std::max(max_ts, (double)ind.tree_size);

    for (auto& ind : pop) {
        if (ind.rank != 1) continue;
        double f1 = ind.mse_domain;
        double f2 = ind.mse_boundary;
        double f3 = (double)ind.tree_size / max_ts;   // normalizar a [0,1]

        // Filtrar puntos fuera de la caja de referencia
        if (f1 >= ref_dom || f2 >= ref_bnd || f3 >= ref_size) continue;
        pts.push_back({f1, f2, f3});
    }

    if (pts.empty()) return 0.0;

    // ── Slice-by-slice sobre f3 (ascendente = menor tree_size primero) ─────
    // Ordenar por f3 ascendente
    std::sort(pts.begin(), pts.end(),
              [](const Pt3& a, const Pt3& b){ return a.f3 < b.f3; });

    double hv3 = 0.0;
    double prev_f3 = pts[0].f3;

    // Para cada "slice" entre f3[i-1] y f3[i], calcular el HV 2D del
    // conjunto acumulado de puntos proyectados sobre (f1, f2).
    std::vector<std::pair<double,double>> slice_pts;

    for (size_t i = 0; i < pts.size(); ++i) {
        slice_pts.push_back({pts[i].f1, pts[i].f2});

        // Grosor del slice en la dimensión f3
        double next_f3 = (i + 1 < pts.size()) ? pts[i + 1].f3 : ref_size;
        double thickness = next_f3 - prev_f3;

        if (thickness > 0.0) {
            double area = hv2d(slice_pts, ref_dom, ref_bnd);
            hv3 += area * thickness;
        }
        prev_f3 = next_f3;
    }

    // Normalizar por el volumen total de la caja de referencia
    return hv3 / (ref_dom * ref_bnd * ref_size);
}

template<typename Ind>
Stats compute_stats(const std::vector<Ind>& pop,
                    const std::string& method,
                    const std::string& pde,
                    double rt)
{
    Stats s;
    s.method = method; s.pde = pde; s.runtime_s = rt;
    for (auto& ind : pop) {
        if (ind.rank == 1) {
            s.front_size++;
            s.best_domain = std::min(s.best_domain, ind.mse_domain);
            s.best_bnd    = std::min(s.best_bnd,    ind.mse_boundary);
            s.mean_domain += ind.mse_domain;
            s.mean_bnd    += ind.mse_boundary;
        }
    }
    if (s.front_size > 0) {
        s.mean_domain /= s.front_size;
        s.mean_bnd    /= s.front_size;
    }
    s.hypervolume = compute_hypervolume(pop);
    return s;
}

// ─── Guardar resumen comparativo ─────────────────────────────────────────────
void save_summary(const std::vector<Stats>& stats, const std::string& path) {
    std::ofstream f(path);
    f << std::fixed << std::setprecision(10);
    f << "method,pde,pareto_size,best_mse_domain,best_mse_boundary,"
         "mean_mse_domain,mean_mse_boundary,hypervolume,runtime_s\n";
    for (auto& s : stats) {
        f << s.method << "," << s.pde << "," << s.front_size << ","
          << s.best_domain << "," << s.best_bnd << ","
          << s.mean_domain << "," << s.mean_bnd << ","
          << s.hypervolume << "," << s.runtime_s << "\n";
    }
}

// ─── Tabla en consola (Solo PI-NSGA-II) ──────────────────────────────────────
void print_table(const std::string& lbl, const Stats& p) {
    std::cout << "\n+------------------+-------------+-------------+----------+-------------+----------+\n";
    std::cout << "| " << std::left << std::setw(78) << ("  Ecuacion: " + lbl) << "|\n";
    std::cout << "+------------------+-------------+-------------+----------+-------------+----------+\n";
    std::cout << "| Metodo           | MSE Dom.    | MSE Bnd.    | Pareto   | Hipervolumen| Tiempo   |\n";
    std::cout << "+------------------+-------------+-------------+----------+-------------+----------+\n";
    std::cout << std::fixed << std::setprecision(4);
    std::cout << "| PI-NSGA-II       | " << std::setw(11) << p.best_domain
              << " | " << std::setw(11) << p.best_bnd
              << " | " << std::setw(8)  << p.front_size
              << " | " << std::setw(11) << p.hypervolume
              << " | " << std::setw(7)  << p.runtime_s << "s |\n";
    std::cout << "+------------------+-------------+-------------+----------+-------------+----------+\n";
}

// ─── Una corrida completa (Solo PI-NSGA-II) ──────────────────────────────────
std::vector<Stats> run_once(int run_id, const std::string& out_dir, bool verbose, bool is_test) {
    unsigned seed_base = 1000u * (unsigned)(run_id + 1);
    std::vector<PDEProblem> problems;
    for (int d : {1, 2}) {
        problems.push_back(make_laplace(d));
        problems.push_back(make_poisson(d));
        problems.push_back(make_helmholtz(d, 1.0));
        problems.push_back(make_schrodinger(d));
        problems.push_back(make_harmonic_oscillator(d));
        problems.push_back(make_airy(d));
        problems.push_back(make_fisher(d));
        problems.push_back(make_duffing(d));
        problems.push_back(make_thomas_fermi(d));
    }

    problems.push_back(make_nonlinear_poisson());
    problems.push_back(make_liouville());
    problems.push_back(make_sine_gordon());
    problems.push_back(make_navier_stokes());

    std::vector<Stats> all_stats;

    for (auto& prob : problems) {
        std::string lbl = prob.name() + (prob.dim == 1 ? "_1D" : "_2D");
        
        if (verbose) std::cout << "\n[Run] PI-NSGA-II en " << lbl << "\n";
        if (prob.is_numerical) {
            prob.numerical_truth = NumericalSolver::solve(prob, 50);
        }

        // Run Numerical Solver (RK4/FDM) baseline for comparison
        bool has_numerical = false;
        double num_rt = 0.0;
        double num_mse_dom = 0.0;
        double num_mse_bnd = 0.0;

        if (prob.dim == 1 || prob.type == PDE::LAPLACE || prob.type == PDE::POISSON || prob.type == PDE::HELMHOLTZ) {
            has_numerical = true;
            auto t_num0 = std::chrono::steady_clock::now();
            auto num_sol = NumericalSolver::solve(prob, 50);
            num_rt = std::chrono::duration<double>(std::chrono::steady_clock::now() - t_num0).count();

            if (prob.is_numerical) {
                // For numerical equations, the numerical solution is the ground truth
                num_mse_dom = 0.0;
                num_mse_bnd = 0.0;
            } else {
                double sum_sq_dom = 0.0;
                if (prob.dim == 1) {
                    int N = num_sol.size();
                    double h = 1.0 / (N - 1);
                    for (int i = 0; i < N; ++i) {
                        double x = i * h;
                        double diff = std::abs(num_sol[i] - prob.exact(x, 0.0));
                        sum_sq_dom += diff * diff;
                    }
                    num_mse_dom = sum_sq_dom / N;
                    num_mse_bnd = 0.0;
                } else {
                    int N = std::sqrt(num_sol.size());
                    double h = 1.0 / (N - 1);
                    for (int i = 0; i < N; ++i) {
                        for (int j = 0; j < N; ++j) {
                            double x = i * h;
                            double y = j * h;
                            double diff = std::abs(num_sol[i * N + j] - prob.exact(x, y));
                            sum_sq_dom += diff * diff;
                        }
                    }
                    num_mse_dom = sum_sq_dom / (N * N);
                    num_mse_bnd = 0.0;
                }
            }
        }

        {
            auto t0 = std::chrono::steady_clock::now();
            PISolver pi(prob, seed_base + 500u);
            int pop = is_test ? 10 : Config::POP_SIZE;
            int gen = is_test ? 5 : Config::MAX_GEN;
            auto pi_pop = pi.run(pop, gen);
            double pi_rt = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
            
            save_pareto_csv(pi_pop, out_dir + "/" + lbl + "_pi_gn_pareto.csv", "PI-NSGA-II", prob.name(), prob.dim);
            save_convergence_csv(pi.history(), out_dir + "/" + lbl + "_pi_gn_convergence.csv");
            save_best_expression(pi_pop, out_dir + "/expr_" + lbl + "_PI-NSGA-II.tex");
            save_best_grid(pi_pop, prob, out_dir + "/grid_" + lbl + "_PI-NSGA-II.csv");
            
            Stats ps = compute_stats(pi_pop, "PI-NSGA-II", lbl, pi_rt);
            if (verbose) print_table(lbl, ps);
            all_stats.push_back(ps);
        }

        if (has_numerical) {
            Stats ns;
            ns.method = "RK4/FDM";
            ns.pde = lbl;
            ns.front_size = 1;
            ns.best_domain = num_mse_dom;
            ns.best_bnd = num_mse_bnd;
            ns.mean_domain = num_mse_dom;
            ns.mean_bnd = num_mse_bnd;
            ns.runtime_s = num_rt;
            ns.hypervolume = 0.0;
            all_stats.push_back(ns);
        }
    }
    return all_stats;
}

int main(int argc, char* argv[]) {
    int n_runs = 1;
    bool is_test = false;
    for (int i = 1; i < argc; ++i) {
        if (std::strcmp(argv[i], "--runs") == 0 && i+1 < argc) n_runs = std::atoi(argv[i+1]);
        if (std::strcmp(argv[i], "--test") == 0) is_test = true;
    }

    if (is_test) {
        int max_threads = omp_get_max_threads();
        int threads_to_use = std::max(1, max_threads - 2);
        omp_set_num_threads(threads_to_use);
        std::cout << "[INFO] Modo --test activado. Usando " << threads_to_use << " nucleos.\n";
    }

    std::cout << "=============================================================\n";
    std::cout << "  PI-NSGA-II --- Orquestador de Ecuaciones PDE\n";
    std::cout << "  Pop=" << Config::POP_SIZE << "  Gen=" << Config::MAX_GEN << "  Runs=" << n_runs << "\n";
    std::cout << "=============================================================\n\n";

    fs::create_directories("results");
    bool verbose = true; 
    std::vector<std::vector<Stats>> all_runs;

    auto t0_all = std::chrono::steady_clock::now();
    for (int r = 0; r < n_runs; ++r) {
        if (n_runs > 1) {
            std::cout << "\n\033[1;34m" << "#############################################################" << "\033[0m\n";
            std::cout << "\033[1;34m" << "  INICIANDO RUN " << r + 1 << " / " << n_runs << "\033[0m\n";
            std::cout << "\033[1;34m" << "#############################################################" << "\033[0m\n";
        }
        
        auto t0_run = std::chrono::steady_clock::now();
        std::string out_dir = (n_runs == 1) ? "results" : "results/run_" + std::to_string(r);
        if (n_runs > 1) fs::create_directories(out_dir);
        
        auto stats = run_once(r, out_dir, verbose, is_test);
        all_runs.push_back(stats);
        save_summary(stats, out_dir + "/comparison_summary.csv");

        if (!verbose) {
            std::cout << "  Run " << r << " completed.\n";
        }
    }

    std::ofstream f_all("results/all_runs_summary.csv");
    f_all << "run,method,pde,pareto_size,best_mse_domain,best_mse_boundary,mean_mse_domain,mean_mse_boundary,hypervolume,runtime_s\n";
    for (int r = 0; r < n_runs; ++r) {
        for (auto& s : all_runs[r]) {
            f_all << r << "," << s.method << "," << s.pde << "," << s.front_size << "," << s.best_domain << ","
                  << s.best_bnd << "," << s.mean_domain << "," << s.mean_bnd << "," << s.hypervolume << "," << s.runtime_s << "\n";
        }
    }

    auto t_total = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0_all).count();
    std::cout << "\n  Benchmark completado en " << t_total << "s. Resultados en ./results/\n";
    return 0;
}