#include "FluidSim/Junction.hpp"
#include "FluidSim/Solver.hpp"
#include <array>
#include <cmath>
#include <complex>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <iomanip>
#include <string>
using namespace fluidsim::port;

namespace {
constexpr double pressure=101325, area=.001;
struct Result { double incident{}, reflected{}, transmitted{}, mass_drift{}, energy_drift{};
    std::array<std::complex<double>,2> incident_spectrum{}, reflected_spectrum{}, transmitted_spectrum{};
    std::uint64_t repairs{}, limited{}, steps{}, fallbacks{}; };
Result run(const std::string& kind,int n,float courant,PipeSolver scheme,
           double volume,float discharge,const std::filesystem::path& trace) {
    const bool straight=kind=="straight", split=kind=="split", star=kind=="star" || kind=="legacy-star" || kind=="star-unequal";
    const int branches=star || kind=="volume-y" ? 3 : 2;
    Junction chamber(float(volume),pressure,300);
    const double rho=chamber.density(), c=std::sqrt(1.4*pressure/rho), z=rho*c;
    std::vector<Pipe1D> pipes;
    if(straight) pipes.emplace_back(area,4,2*n);
    else for(int j=0;j<branches;++j) pipes.emplace_back(area*(kind=="star-unequal"?j+1:1),2,n);
    for(auto& p:pipes) {
        p.solver=scheme; p.damping_multiplier=0; p.energy_relaxation_rate=0;
        for(int i=0;i<int(p.cells.size());++i) {
            const double x=(i+.5)*p.dx;
            const double pulse=&p==&pipes[0] ? 100*std::exp(-std::pow((x-.6)/.12,2)) : 0;
            const double density=rho+pulse/(c*c), u=pulse/z;
            p.cells[i]={float(density),float(density*u),float((pressure+pulse)/pipe_gm1+.5*density*u*u),1,0};
        }
    }
    BoundarySystem boundaries(pipes);
    boundaries.add_closed_end(0,true);
    if(straight) boundaries.add_closed_end(0,false);
    else {
        for(int j=1;j<branches;++j) boundaries.add_closed_end(j,false);
        if(split) boundaries.add_connection(0,false,1,true);
        else if(star) boundaries.add_junction({{0,false},{1,true},{2,true}},kind=="legacy-star"?BoundarySystem::JunctionMethod::LegacyRusanov:BoundarySystem::JunctionMethod::Characteristic);
        else {
            std::vector<JunctionEnd> ends{{0,false,discharge},{1,true,discharge}};
            if(branches==3) ends.push_back({2,true,discharge});
            chamber.connect(boundaries,pipes,ends);
        }
    }
    Solver solver(pipes,boundaries); solver.minimum_substeps=1;
    const bool has_volume=!straight && !split && !star;
    if(has_volume) solver.add_component(chamber);
    auto totals=[&] {
        std::array<double,2> t{has_volume?chamber.mass():0,has_volume?chamber.internal_energy():0};
        for(const auto& p:pipes) for(int i=0;i<int(p.cells.size());++i) {
            t[0]+=p.cell_mass(i); t[1]+=double(p.cells[i].energy)*p.cell_area(i)*p.cell_length(i);
        }
        return t;
    };
    const auto initial=totals(); Result r;
    std::ofstream out(trace);
    if(!out) throw std::runtime_error("cannot open trace");
    out<<std::setprecision(10);
    out<<"time_s,upstream_pressure_pa,upstream_velocity_m_s,downstream_pressure_pa,downstream_velocity_m_s\n";
    const int up=int(1.4/pipes[0].dx);
    const int down=int((straight?2.6:.6)/pipes.back().dx);
    double time=0;
    // Stop before far-end reflections return to either probe. Integrate acoustic
    // characteristics p+/- = (p' +/- rho*c*u)/2, not squared total pressure.
    while(time<.009) {
        const float dt=float(std::min(.009-time,courant*(2./n)/(c+.5)));
        const auto a=pipes[0].state(up), b=pipes.back().state(down);
        const double plus=(a.pressure-pressure+z*a.velocity)/2;
        const double minus=(a.pressure-pressure-z*a.velocity)/2;
        for(int frequency=0;frequency<2;++frequency) {
            const double hz=frequency==0?500:1000;
            const auto weight=std::exp(std::complex<double>(0,-2*3.141592653589793*hz*time))*double(dt);
            r.incident_spectrum[frequency]+=plus*weight;
            r.reflected_spectrum[frequency]+=minus*weight;
            r.transmitted_spectrum[frequency]+=(b.pressure-pressure+z*b.velocity)*.5*weight;
        }
        r.incident+=plus*plus*area/z*dt;
        r.reflected+=minus*minus*area/z*dt;
        if(straight) {
            const double v=(b.pressure-pressure+z*b.velocity)/2;
            r.transmitted+=v*v*area/z*dt;
        } else for(int j=1;j<branches;++j) {
            const auto q=pipes[j].state(down);
            const double v=(q.pressure-pressure+z*q.velocity)/2;
            r.transmitted+=v*v*pipes[j].cell_area(down)/z*dt;
        }
        out<<time<<','<<a.pressure-pressure<<','<<a.velocity<<','<<b.pressure-pressure<<','<<b.velocity<<'\n';
        solver.step(dt); time+=dt;
    }
    const auto final=totals();
    r.mass_drift=(final[0]-initial[0])/initial[0];
    r.energy_drift=(final[1]-initial[1])/initial[1];
    for(auto& p:pipes) r.repairs+=p.positivity_repairs;
    r.fallbacks=boundaries.junction_fallbacks;
    r.limited=boundaries.limited_fluxes; r.steps=solver.substep_count();
    if(!(r.incident>0) || !std::isfinite(r.transmitted) || r.repairs ||
       std::abs(r.mass_drift)>1e-4 || std::abs(r.energy_drift)>1e-4)
        throw std::runtime_error("acoustic benchmark robustness/conservation failure: "+trace.string()+" mass="+std::to_string(r.mass_drift)+" energy="+std::to_string(r.energy_drift)+" limited="+std::to_string(r.limited)+" repairs="+std::to_string(r.repairs));
    return r;
}
}
int main(int argc,char** argv) {
    try {
        const std::filesystem::path dir=argc>1?argv[1]:"out/acoustics";
        std::filesystem::create_directories(dir);
        std::ofstream summary(dir/"summary.csv");
        summary<<std::setprecision(10);
        summary<<"case,solver,cells_per_2m,courant,volume_m3,discharge,incident_j,reflected_ratio,transmitted_ratio,unaccounted_ratio,mass_relative_drift,energy_relative_drift,limited_fluxes,positivity_repairs,substeps,transmission_over_straight,reflection_over_straight,ideal_reflected_ratio,ideal_transmitted_ratio,status,junction_fallbacks,t500_gain_over_straight,t500_phase_error_rad,t1000_gain_over_straight,t1000_phase_error_rad\n";
        double straight_transmission=0, straight_reflection=0, coarse_transmission=0;
        std::array<double,6> first_order_transmission{};
        std::array<std::complex<double>,2> straight_spectrum{};
        int cases=0, flagged=0;
        for(const auto* scheme:{"rusanov","hlle","hllc","muscl"})
        for(int n:{40,80,160}) for(float courant:{.4f,.8f})
        for(const auto* kind:{"straight","split","star","legacy-star","star-unequal","volume","volume-y"})
        for(double volume:{.00005,.0005,.005}) {
            if((std::string(kind)=="straight" || std::string(kind)=="split" || (std::string(kind)=="star" || std::string(kind)=="legacy-star" || std::string(kind)=="star-unequal")) && volume!=.00005) continue;
            for(float discharge:{1.f,.7f}) {
                if((std::string(kind)=="straight" || std::string(kind)=="split" || (std::string(kind)=="star" || std::string(kind)=="legacy-star" || std::string(kind)=="star-unequal")) && discharge!=1) continue;
                const std::string id=std::string(kind)+"-"+scheme+"-"+std::to_string(n)+"-"+std::to_string(courant)+"-"+std::to_string(volume)+"-"+std::to_string(discharge);
                const auto r=run(kind,n,courant,parse_pipe_solver(scheme),volume,discharge,dir/(id+".csv"));
                const double reflected=r.reflected/r.incident, transmitted=r.transmitted/r.incident;
                if(std::string(kind)=="straight") {
                    straight_transmission=transmitted; straight_reflection=reflected;
                    for(int frequency=0;frequency<2;++frequency)
                        straight_spectrum[frequency]=r.transmitted_spectrum[frequency]/r.incident_spectrum[frequency];
                    const int slot=(n==40?0:n==80?2:4)+(courant>.5f?1:0);
                    if(std::string(scheme)=="rusanov") first_order_transmission[slot]=transmitted;
                    if(std::string(scheme)=="muscl" && transmitted<first_order_transmission[slot]+.05)
                        throw std::runtime_error("MUSCL pulse attenuation improvement regressed");
                    if(reflected>1e-4 || transmitted<=0 || transmitted>1.01)
                        throw std::runtime_error("straight acoustic propagation regression");
                    if(courant==.4f) {
                        if(n>40 && transmitted<=coarse_transmission)
                            throw std::runtime_error("straight pulse does not converge with mesh refinement");
                        coarse_transmission=transmitted;
                    }
                }
                if(std::string(kind)=="split" &&
                   (std::abs(transmitted-straight_transmission)>1e-7 || std::abs(reflected-straight_reflection)>1e-7))
                    throw std::runtime_error("split connection differs from continuous pipe");
                if(std::string(kind)=="star" || std::string(kind)=="star-unequal") {
                    const double expected_r=std::string(kind)=="star"?1./9:4./9;
                    const double expected_t=1-expected_r;
                    if(std::abs(reflected/straight_transmission-expected_r)>.04 ||
                       std::abs(transmitted/straight_transmission-expected_t)>.07 || r.fallbacks)
                        throw std::runtime_error("characteristic junction acoustic scattering regression: "+id+" R="+std::to_string(reflected/straight_transmission)+" T="+std::to_string(transmitted/straight_transmission)+" fallback="+std::to_string(r.fallbacks));
                }
                const bool contaminated=r.limited || reflected+transmitted>1.05;
                ++cases; flagged+=contaminated;
                summary<<kind<<','<<scheme<<','<<n<<','<<courant<<','<<volume<<','<<discharge<<','<<r.incident<<','<<reflected<<','<<transmitted<<','<<1-reflected-transmitted<<','<<r.mass_drift<<','<<r.energy_drift<<','<<r.limited<<','<<r.repairs<<','<<r.steps<<','<<transmitted/straight_transmission<<','<<reflected/straight_transmission<<',';
                if(std::string(kind)=="straight" || std::string(kind)=="split") summary<<"0,1,";
                else if(std::string(kind)=="star-unequal") summary<<4./9<<','<<5./9<<',';
                else if(std::string(kind)=="star" || std::string(kind)=="legacy-star") summary<<1./9<<','<<8./9<<',';
                else summary<<",,";
                summary<<(contaminated?"flagged":"ok")<<','<<r.fallbacks;
                for(int frequency=0;frequency<2;++frequency) {
                    const auto corrected=r.transmitted_spectrum[frequency]/r.incident_spectrum[frequency]/straight_spectrum[frequency];
                    summary<<','<<std::abs(corrected)<<','<<std::arg(corrected);
                    if(std::string(kind)=="star" && n==160 && (std::abs(std::abs(corrected)-2./3)>.08 || std::abs(std::arg(corrected))>.15))
                        throw std::runtime_error("junction frequency/phase regression: "+id);
                }
                summary<<'\n';
            }
        }
        std::cout<<"Validated "<<cases<<" cases ("<<flagged<<" flagged for limiting or acoustic amplification). Wrote acoustic pulse benchmark to "<<dir<<'\n';
    } catch(const std::exception& e) { std::cerr<<e.what()<<'\n'; return 1; }
}