#include "FluidSim/Turbocharger.hpp"
#include <algorithm>
#include <cmath>
#include <stdexcept>

namespace fluidsim::port {
namespace {
constexpr double cp = double(gamma)*gas_constant/(gamma-1);
constexpr double exponent = double(gamma-1)/gamma;
TurboConfig checked(TurboConfig c) {
    for (double x : {c.shaft_inertia, c.compressor_radius, c.turbine_radius,
                     c.compressor_head_coefficient, c.turbine_optimum_speed_ratio,
                     c.compressor_effective_area, c.turbine_effective_area, c.gate_diaphragm_area,
                     c.gate_spring_stiffness, c.gate_damping, c.gate_diameter, c.reference_pressure})
        if (!std::isfinite(x) || x <= 0) throw std::invalid_argument("invalid turbo geometry/inertia");
    for (double x : {c.compressor_efficiency, c.turbine_peak_efficiency})
        if (!std::isfinite(x) || x <= 0 || x > 1) throw std::invalid_argument("invalid turbo efficiency");
    for (double x : {c.bearing_viscous, c.windage, c.gate_spring_preload, c.gate_maximum_lift})
        if (!std::isfinite(x) || x < 0) throw std::invalid_argument("invalid turbo loss coefficient");
    for (float x : c.connection_diameters)
        if (!std::isfinite(x) || x <= 0) throw std::invalid_argument("invalid turbo connection diameter");
    return c;
}
}
Turbocharger::Turbocharger(TurboConfig configuration)
    : config(checked(configuration)),
      compressor_inlet(config.compressor_inlet_volume), compressor_outlet(config.compressor_outlet_volume),
      turbine_inlet(config.turbine_inlet_volume), turbine_outlet(config.turbine_outlet_volume),
      compressor_inlet_port(compressor_inlet.create_port()),
      compressor_outlet_port(compressor_outlet.create_port()),
      turbine_inlet_port(turbine_inlet.create_port()), turbine_outlet_port(turbine_outlet.create_port()),
      internal_{&compressor_inlet.create_port(), &compressor_outlet.create_port(),
                &turbine_inlet.create_port(), &turbine_outlet.create_port()} {}

void Turbocharger::connect(BoundarySystem& b, const std::vector<Pipe1D>& pipes,
                          int ci, bool cil, int co, bool col, int ti, bool til, int to, bool tol) {
    if (connected_) throw std::logic_error("turbo already connected");
    const std::array<int, 4> indices{ci, co, ti, to};
    const std::array<bool, 4> left{cil, col, til, tol};
    auto external = ports();
    for (std::size_t i = 0; i < 4; ++i) {
        const auto& p = pipes.at(indices[i]);
        const float diameter = config.connection_diameters[i];
        connection_areas_[i] = std::min(p.cell_area(left[i] ? 0 : int(p.cells.size())-1), pi*diameter*diameter*.25f);
        b.add_orifice(*external[i], indices[i], left[i], connection_areas_[i]);
    }
    connected_ = true;
}
std::vector<Port*> Turbocharger::ports() {
    return {&compressor_inlet_port, &compressor_outlet_port, &turbine_inlet_port, &turbine_outlet_port};
}
double Turbocharger::fluid_energy() const {
    return double(compressor_inlet.internal_energy()) + compressor_outlet.internal_energy() +
           turbine_inlet.internal_energy() + turbine_outlet.internal_energy();
}
void Turbocharger::set_initial_rpm(double rpm) {
    if (!std::isfinite(rpm) || rpm < 0) throw std::invalid_argument("invalid initial turbo speed");
    omega_ = rpm*2*pi/60;
}
double Turbocharger::nozzle(double p, double t, double back, double area) {
    if (back >= p) return 0;
    const double critical = std::pow(2/(gamma+1), gamma/(gamma-1));
    const double ratio = std::max(back/p, critical);
    const double factor = 2*gamma/(gas_constant*(gamma-1)) *
                          (std::pow(ratio, 2/gamma)-std::pow(ratio, (gamma+1)/gamma));
    return area*p/std::sqrt(t)*std::sqrt(std::max(0., factor));
}
double Turbocharger::transfer_flow(Volume0D& in, Volume0D& out, Port& in_ext,
                                  Port& out_ext, double requested, double dt) {
    const bool forward = requested >= 0;
    const auto& donor = forward ? in : out;
    const auto& external = forward ? in_ext : out_ext;
    const double h = cp*donor.temperature();
    // Reserve the mass/energy already promised to the external pipe boundary.
    const double mass = donor.mass()+std::min(0.f, external.mass_flow)*dt;
    const double energy = donor.internal_energy()+std::min(0.f, external.mass_flow)*external.enthalpy*dt;
    const double bound = .25*std::max(0., std::min(mass, energy/h))/dt;
    const double actual = std::clamp(requested, -bound, bound);
    if (actual != requested) ++limited_transfers_;
    return actual;
}
void Turbocharger::set_transfer(Volume0D& in, Volume0D& out, Port& ip, Port& op,
                               double flow, double work, double dt) {
    const auto& donor = flow >= 0 ? in : out;
    const double h = cp*donor.temperature();
    ip.mass_flow = float(-flow);
    op.mass_flow = float(flow);
    ip.air_fraction = op.air_fraction = donor.air_fraction();
    ip.enthalpy = op.enthalpy = float(h);
    // Positive work means energy delivered to the shaft. Only the receiving
    // fluid enthalpy changes; the donor loses precisely its stagnation enthalpy.
    if (std::abs(flow) > 0) {
        auto& receiver = flow >= 0 ? op : ip;
        receiver.enthalpy = float(h-work/(std::abs(flow)*dt));
    }
}
void Turbocharger::update_state(float timestep) {
    if (!std::isfinite(timestep) || timestep < 0) throw std::invalid_argument("invalid turbo timestep");
    if (timestep == 0) return;
    const double dt = timestep, initial_fluid = fluid_energy(), initial_shaft = shaft_energy();
    const std::array<const Port*, 4> external{&compressor_inlet_port, &compressor_outlet_port,
                                             &turbine_inlet_port, &turbine_outlet_port};
    double external_work = 0;
    for (const auto* p : external) external_work += double(p->mass_flow)*p->enthalpy*dt;

    const double pin = compressor_inlet.pressure(), tin = compressor_inlet.temperature();
    const double head = config.compressor_head_coefficient*std::pow(omega_*config.compressor_radius, 2);
    const double pump_pressure = pin*std::pow(1+config.compressor_efficiency*head/(cp*tin), 1/exponent);
    double mc = compressor_outlet.pressure() <= pump_pressure
        ? nozzle(pump_pressure, tin+head/cp, compressor_outlet.pressure(), config.compressor_effective_area)
        : -nozzle(compressor_outlet.pressure(), compressor_outlet.temperature(), pump_pressure,
                  config.compressor_effective_area);
    mc = transfer_flow(compressor_inlet, compressor_outlet, compressor_inlet_port, compressor_outlet_port, mc, dt);

    const double tp = turbine_inlet.pressure(), tout = turbine_outlet.pressure();
    const double can_force = std::max(0., double(compressor_outlet.pressure())-config.reference_pressure)*
                             config.gate_diaphragm_area;
    const double equilibrium_lift = (can_force-config.gate_spring_preload)/config.gate_spring_stiffness;
    // Overdamped spring-can linkage: b*dx/dt + k*x = F-Fpreload, with seats.
    wastegate_lift_ = std::clamp(equilibrium_lift+(wastegate_lift_-equilibrium_lift)*
        std::exp(-dt*config.gate_spring_stiffness/config.gate_damping), 0., config.gate_maximum_lift);
    const double gate_area = pi*config.gate_diameter*wastegate_lift_;
    double mt = tp >= tout ? nozzle(tp, turbine_inlet.temperature(), tout, config.turbine_effective_area)
        : -nozzle(tout, turbine_outlet.temperature(), tp, config.turbine_effective_area);
    double mw = tp >= tout ? nozzle(tp, turbine_inlet.temperature(), tout, gate_area)
        : -nozzle(tout, turbine_outlet.temperature(), tp, gate_area);
    const double requested = mt+mw;
    const double actual = transfer_flow(turbine_inlet, turbine_outlet, turbine_inlet_port, turbine_outlet_port, requested, dt);
    if (requested != 0) { mt *= actual/requested; mw *= actual/requested; }
    const double ideal_drop = tp > tout ? cp*turbine_inlet.temperature()*(1-std::pow(tout/tp, exponent)) : 0;
    const double jet = std::sqrt(2*ideal_drop), radius = config.turbine_radius;
    const double ratio = config.turbine_optimum_speed_ratio, eta = config.turbine_peak_efficiency;
    const double drive = std::max(0., mt)*eta*jet*radius/ratio;
    const double turbine_drag = std::max(0., mt)*eta*.5*radius*radius/(ratio*ratio);
    const double compressor_drag = std::max(0., mc)*config.compressor_head_coefficient*
                                   config.compressor_radius*config.compressor_radius;
    // Implicit midpoint solves I*domega/dt = drive - b*omega - c*omega^2.
    // Evaluating work at midpoint conserves rotor kinetic energy, even at rest.
    double wt = 0, wc = 0, loss = 0, remaining = dt;
    int steps = 0;
    while (remaining > dt*1e-12) {
        const double drag = turbine_drag+compressor_drag+config.bearing_viscous+config.windage*omega_;
        const double h = std::min(remaining, config.shaft_inertia/std::max(drag, 1e-300));
        if (!(h > 0) || ++steps > 1000000) throw std::runtime_error("turbo shaft timestep failed");
        const double b = 2*config.shaft_inertia/h+turbine_drag+compressor_drag+config.bearing_viscous;
        const double rhs = 2*config.shaft_inertia/h*omega_+drive;
        const double mid = 2*rhs/(b+std::sqrt(b*b+4*config.windage*rhs));
        const double next = 2*mid-omega_;
        if (next < 0) throw std::runtime_error("turbo shaft integration reversed unexpectedly");
        wt += (drive-turbine_drag*mid)*mid*h;
        wc += compressor_drag*mid*mid*h;
        loss += (config.bearing_viscous+config.windage*mid)*mid*mid*h;
        omega_ = next;
        angle_ = std::fmod(angle_+mid*h, 2*double(pi));
        remaining -= h;
    }
    bearing_heat_ += loss;
    turbine_power_ = wt/dt;
    compressor_power_ = wc/dt;
    bearing_power_ = loss/dt;
    turbine_flow_ = mt;
    wastegate_flow_ = mw;
    compressor_flow_ = mc;
    set_transfer(compressor_inlet, compressor_outlet, *internal_[0], *internal_[1], mc, -wc, dt);
    set_transfer(turbine_inlet, turbine_outlet, *internal_[2], *internal_[3], mt+mw, wt, dt);
    compressor_inlet.update_state(timestep);
    compressor_outlet.update_state(timestep);
    turbine_inlet.update_state(timestep);
    turbine_outlet.update_state(timestep);
    energy_residual_ += fluid_energy()-initial_fluid + shaft_energy()-initial_shaft+loss-external_work;
}
TurboSnapshot Turbocharger::snapshot() const {
    return {omega_*60/(2*pi), angle_, shaft_energy(), turbine_power_, compressor_power_, bearing_power_,
            compressor_flow_, turbine_flow_, compressor_outlet.pressure()/compressor_inlet.pressure(),
            turbine_inlet.pressure()/turbine_outlet.pressure(), turbine_inlet.pressure(),
            turbine_outlet.pressure(), compressor_outlet.pressure(), compressor_outlet.temperature(),
            bearing_heat_, energy_residual_, limited_transfers_, wastegate_flow_, wastegate_lift_};
}
} // namespace fluidsim::port