#include "Rocket.h"
#include "Stage.h"
#include "Environment.h"
#include "FlightState.h"
#include "Airframe.h"
#include "FinSet.h"
#include "RecoverySystem.h"
#include "Motor.h"

#include <numeric> // for std::accumulate

// Constructor
Rocket::Rocket(const std::string& name)
    : name_(name),
      totalMass_(0.0),
      totalPropellantMass_(0.0),
      centerOfGravity_(0.0),
      centerOfPressure_(0.0)
{}

// Public Methods

double Rocket::getTotalMass() const {
    return totalMass_;
}

double Rocket::getTotalPropellantMass() const {
    return totalPropellantMass_;
}

double Rocket::getCenterOfGravity() const {
    return centerOfGravity_;
}

double Rocket::getCenterOfPressure() const {
    return centerOfPressure_;
}

double Rocket::getStabilityMargin() const {
    // Positive margin (CG ahead of CP) indicates stability.
    // Normalize by diameter or calibers if needed later.
    return centerOfGravity_ - centerOfPressure_;
}

void Rocket::addStage(std::unique_ptr<Stage> stage) {
    stages_.emplace_back(std::move(stage));
}

const std::vector<std::unique_ptr<Stage>>& Rocket::getStages() const {
    return stages_;
}

void Rocket::updateMassProperties() {
    totalMass_ = 0.0;
    totalPropellantMass_ = 0.0;
    centerOfGravity_ = 0.0;
    centerOfPressure_ = 0.0;

    if (stages_.empty()) {
        return;
    }

    double weightedCGSum = 0.0;
    double weightedCPSum = 0.0;
    double currentZ = 0.0; // Assume stacking stages vertically

    for (const auto& stage : stages_) {
        stage->updateMassProperties();

        double stageMass = stage->getTotalMass();
        double stagePropellantMass = stage->getTotalPropellantMass();

        totalMass_ += stageMass;
        totalPropellantMass_ += stagePropellantMass;

        // For center of gravity, weight by mass
        weightedCGSum += (currentZ + stage->getAirframeLength() / 2.0) * stageMass;

        // For center of pressure, weight by aerodynamic properties
        weightedCPSum += (currentZ + stage->getAirframeLength() / 2.0) * stage->calculateNormalForceCoefficient();

        currentZ += stage->getAirframeLength(); // Move up for next stage
    }

    centerOfGravity_ = weightedCGSum / totalMass_;

    // Note: CP calculation here is a rough average based on fin normal forces
    centerOfPressure_ = weightedCPSum / 
        std::accumulate(stages_.begin(), stages_.end(), 0.0,
                        [](double sum, const std::unique_ptr<Stage>& s) {
                            return sum + s->calculateNormalForceCoefficient();
                        });
}

void Rocket::prepareForFlight(const Environment& env) {
    for (auto& stage : stages_) {
        stage->prepareForFlight(env);
    }
}

void Rocket::applyFlightState(const FlightState& state) {
    for (auto& stage : stages_) {
        stage->applyFlightState(state);
    }
}

const std::string& Rocket::getName() const {
    return name_;
}

void Rocket::setName(const std::string& name) {
    name_ = name;
}
