#include "FluidSim/Turbocharger.hpp"
#include "FluidSim/Atmosphere.hpp"
#include "FluidSim/Solver.hpp"
#include <iostream>
#include <stdexcept>

using namespace fluidsim::port;
namespace {
void check(bool condition, const char* message) { if (!condition) throw std::runtime_error(message); }
double mass(const Turbocharger& t) {
    return double(t.compressor_inlet.mass())+t.compressor_outlet.mass()+t.turbine_inlet.mass()+t.turbine_outlet.mass();
}
double energy(const Turbocharger& t) { return t.fluid_energy()+t.shaft_energy()+t.snapshot().bearing_heat; }
void closed_balance(double initial_rpm) {
    Turbocharger t;
    t.set_initial_rpm(initial_rpm);
    t.turbine_inlet.set_pressure(250000, 900);
    const double before = energy(t), initial_mass = mass(t);
    for (int i = 0; i < 10000; ++i) t.update_state(1e-6f);
    check(std::abs(mass(t)-initial_mass) < initial_mass*3e-5, "turbo created/lost gas mass");
    check(std::abs(energy(t)-before) < before*2e-4, "fluid + shaft + bearing heat energy drift");
    check(t.omega() > 0, "turbine did not start its shaft from pressure energy");
    const auto s = t.snapshot();
    check(std::abs(s.energy_residual) < before*2e-4, "turbo energy residual excessive");
    std::cout << "Closed balance initial_rpm=" << initial_rpm << " mass_error=" << mass(t)-initial_mass
              << " energy_error=" << energy(t)-before << " J\n";
}
double turbine_flow(float downstream) {
    Turbocharger t;
    t.turbine_inlet.set_pressure(300000, 900);
    t.turbine_outlet.set_pressure(downstream, 600);
    t.update_state(1e-7f);
    return t.snapshot().turbine_flow;
}
}
int main() {
    Turbocharger rest;
    const double e = energy(rest);
    for (int i = 0; i < 1000; ++i) rest.update_state(1.f/44100);
    check(rest.omega() == 0 && energy(rest) == e && rest.snapshot().compressor_pressure_ratio == 1,
          "stationary equilibrium generated boost/shaft energy");
    closed_balance(0);
    closed_balance(120000);
    check(std::abs(turbine_flow(100000)-turbine_flow(50000)) < 1e-10, "choked turbine flow depends on back pressure");
    check(turbine_flow(250000) < turbine_flow(100000), "higher backpressure did not reduce subcritical flow");
    Turbocharger reverse;
    reverse.compressor_outlet.set_pressure(150000, 350);
    reverse.turbine_outlet.set_pressure(150000, 600);
    const auto reverse_energy = energy(reverse);
    reverse.update_state(1e-6f);
    check(reverse.snapshot().compressor_flow < 0 && reverse.snapshot().turbine_flow < 0,
          "turbo acted as an artificial one-way valve");
    check(std::abs(energy(reverse)-reverse_energy) < reverse_energy*1e-6, "reverse flow creates energy");
    Turbocharger gate;
    gate.compressor_outlet.set_pressure(220000, 350);
    gate.turbine_inlet.set_pressure(250000, 900);
    gate.update_state(1e-6f);
    check(gate.snapshot().wastegate_lift > 0 && gate.snapshot().wastegate_flow > 0,
          "pressure can did not open the wastegate bypass");
    check(std::abs(gate.snapshot().energy_residual) < 1e-4, "wastegate bypass created energy");
    Turbocharger coarse;
    coarse.set_initial_rpm(200000);
    const double coarse_energy = energy(coarse);
    coarse.update_state(.5f);
    check(coarse.omega() >= 0 && std::abs(energy(coarse)-coarse_energy) < coarse_energy*1e-5,
          "large turbo timestep creates energy or reverses shaft");
    bool caught = false;
    try { TurboConfig bad; bad.shaft_inertia = 0; Turbocharger invalid(bad); }
    catch (const std::invalid_argument&) { caught = true; }
    check(caught, "zero rotor inertia accepted");

    // Hot pulsing reservoir -> real pipes -> turbine -> exhaust reservoir.
    // Compressor feeds a pressure reservoir representing a loaded charge duct.
    std::vector<Pipe1D> pipes;
    for (int i = 0; i < 4; ++i) pipes.emplace_back(pi*.024f*.024f*.25f, .08f, 8);
    Atmosphere fresh, charge, hot, exhaust;
    charge.pressure = 130000;
    hot.temperature = 900;
    hot.pressure = 220000;
    charge.update_state(0); hot.update_state(0);
    BoundarySystem boundaries(pipes);
    Turbocharger turbo;
    turbo.connect(boundaries, pipes, 0, false, 1, true, 2, false, 3, true);
    std::array<float, 4> areas;
    const std::array<Atmosphere*, 4> reservoirs{&fresh, &charge, &hot, &exhaust};
    for (int i = 0; i < 4; ++i) {
        areas[i] = pipes[i].area;
        boundaries.add_orifice(reservoirs[i]->port, i, i%2 == 0, areas[i]);
    }
    Solver solver(pipes, boundaries);
    solver.add_component(turbo);
    float minimum_power = 1e9f, maximum_power = -1e9f;
    double initial_speed = 0;
    for (int i = 0; i < 44100; ++i) {
        hot.pressure = 220000+60000*std::sin(2*pi*80*i/44100);
        hot.update_state(0);
        solver.step(1.f/44100);
        const auto s = turbo.snapshot();
        if (i == 100) initial_speed = s.rpm;
        if (i > 40000) {
            minimum_power = std::min(minimum_power, float(s.turbine_power));
            maximum_power = std::max(maximum_power, float(s.turbine_power));
        }
    }
    const auto result = turbo.snapshot();
    check(result.rpm > initial_speed+10000, "hot gas did not spool the coupled turbo");
    check(result.compressor_flow > 0 && result.compressor_power > 0 && result.charge_temperature > 300,
          "loaded compressor did not deliver heated charge using shaft work");
    check(maximum_power-minimum_power > 100, "exhaust pulses were averaged away before turbine torque");
    for (const auto& pipe : pipes) check(pipe.positivity_repairs == 0, "turbo bench required pipe repairs");
    check(result.limited_transfers == 0, "normal turbo bench depleted a storage node");
    std::cout << "Turbo pipe bench passed: rpm=" << result.rpm << " compressor flow=" << result.compressor_flow
              << " kg/s, charge T=" << result.charge_temperature << " K, turbine power range="
              << minimum_power << ".." << maximum_power << " W, residual=" << result.energy_residual << " J\n";
}