// src/ForcesModel.cpp

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

#include <cmath>    // For sqrt

// Constructor
ForcesModel::ForcesModel(std::shared_ptr<Rocket> rocket,
                         std::shared_ptr<Environment> environment)
    : rocket_(rocket),
      environment_(environment)
{}

// Public Methods

std::array<double, 3> ForcesModel::computeNetForce(const FlightState& state) const {
    std::array<double, 3> netForce = {0.0, 0.0, 0.0};

    // Thrust
    auto thrust = computeThrust(state);
    for (int i = 0; i < 3; ++i) {
        netForce[i] += thrust[i];
    }

    // Aerodynamic Drag
    auto drag = computeAerodynamicDrag(state);
    for (int i = 0; i < 3; ++i) {
        netForce[i] += drag[i];
    }

    // Gravity
    auto gravity = computeGravity(state);
    for (int i = 0; i < 3; ++i) {
        netForce[i] += gravity[i];
    }

    return netForce;
}

std::array<double, 3> ForcesModel::computeNetMoment(const FlightState& /*state*/) const {
    // Placeholder: No moments in basic 3-DoF model
    return {0.0, 0.0, 0.0};
}

// Private Methods

std::array<double, 3> ForcesModel::computeAerodynamicDrag(const FlightState& state) const {
    std::array<double, 3> dragForce = {0.0, 0.0, 0.0};

    // Get air density at altitude
    double altitude = state.getPosition()[2];
    double airDensity = environment_->getAirDensity(altitude);

    // Get velocity
    const auto& velocity = state.getVelocity();
    double speed = std::sqrt(velocity[0]*velocity[0] +
                             velocity[1]*velocity[1] +
                             velocity[2]*velocity[2]);

    if (speed == 0.0) {
        return dragForce; // No drag if stationary
    }

    // Drag direction is opposite to velocity
    std::array<double, 3> dragDirection = {-velocity[0]/speed,
                                           -velocity[1]/speed,
                                           -velocity[2]/speed};

    // Approximate total reference area and Cd
    double totalReferenceArea = 0.0;
    double effectiveDragCoefficient = 0.0;

    for (const auto& stage : rocket_->getStages()) {
        if (stage->getAirframeLength() > 0.0) {
            totalReferenceArea += stage->getAirframeLength(); // Simple approx
            effectiveDragCoefficient += 0.75; // Assume average Cd
        }
    }

    // Recovery drag if deployed
    for (const auto& stage : rocket_->getStages()) {
        // Assume only one recovery system per stage for now
        // (Can easily extend)
        auto recovery = stage->checkRecoveryEvent(); // Would return a bool now
        if (recovery) {
            totalReferenceArea += 1.0; // Large area
            effectiveDragCoefficient += 1.5; // Large Cd
        }
    }

    if (totalReferenceArea <= 0.0) {
        totalReferenceArea = 0.01; // Prevent division by zero
    }

    // Compute drag force magnitude
    double dragMagnitude = 0.5 * airDensity * speed * speed *
                            effectiveDragCoefficient * totalReferenceArea;

    // Apply direction
    for (int i = 0; i < 3; ++i) {
        dragForce[i] = dragMagnitude * dragDirection[i];
    }

    return dragForce;
}

std::array<double, 3> ForcesModel::computeThrust(const FlightState& /*state*/) const {
    std::array<double, 3> thrustForce = {0.0, 0.0, 0.0};

    for (const auto& stage : rocket_->getStages()) {
        auto stageThrust = stage->getTotalThrust();
        for (int i = 0; i < 3; ++i) {
            thrustForce[i] += stageThrust[i];
        }
    }

    return thrustForce;
}

std::array<double, 3> ForcesModel::computeGravity(const FlightState& state) const {
    std::array<double, 3> gravityForce = {0.0, 0.0, 0.0};

    double altitude = state.getPosition()[2];

    // TODO: Make this more robust. Don't accelerate through the ground
    if(altitude <= 0.0) {
        return {0.0, 0.0, 0.0};
    }

    double g = environment_->getGravity(altitude);

    double totalMass = rocket_->getTotalMass();

    gravityForce[2] = -totalMass * g; // Gravity acts downward

    return gravityForce;
}
