/*
 * Copyright (c) 2025 Centre Tecnologic de Telecomunicacions de Catalunya (CTTC)
 * Author: Amir Ashtari <aashtari@cttc.es>
 * SPDX-License-Identifier: GPL-2.0-only
 */

#include "sionna-rt-channel-model.h"

#include "matrix-based-channel-model.h"
#include "pybind11/numpy.h"
#include "pybind11/pybind11.h"

#include "ns3/antenna-model.h"
#include "ns3/double.h"
#include "ns3/log.h"
#include "ns3/mobility-model.h"
#include "ns3/node.h"
#include "ns3/phased-array-model.h"
#include "ns3/pointer.h"
#include "ns3/simulator.h"
#include "ns3/string.h"
#include "ns3/three-gpp-antenna-model.h"
#include "ns3/uniform-planar-array.h"

#include <filesystem>
namespace py = pybind11;

namespace ns3
{

NS_LOG_COMPONENT_DEFINE("SionnaRtChannelModel");

NS_OBJECT_ENSURE_REGISTERED(SionnaRtChannelModel);

/// Conversion factor: degrees to radians
static constexpr double DEG2RAD = M_PI / 180.0;

SionnaRtChannelModel::SionnaRtChannelModel()
{
    NS_LOG_FUNCTION(this);
}

SionnaRtChannelModel::~SionnaRtChannelModel()
{
    NS_LOG_FUNCTION(this);
}

void
SionnaRtChannelModel::DoDispose()
{
    NS_LOG_FUNCTION(this);
    m_channelMatrixMap.clear();
    m_channelParamsMap.clear();
}

TypeId
SionnaRtChannelModel::GetTypeId()
{
    static TypeId tid =
        TypeId("ns3::SionnaRtChannelModel")
            .SetGroupName("Spectrum")
            .SetParent<MatrixBasedChannelModel>()
            .AddConstructor<SionnaRtChannelModel>()
            .AddAttribute("Frequency",
                          "The operating Frequency in Hz",
                          DoubleValue(500.0e6),
                          MakeDoubleAccessor(&SionnaRtChannelModel::SetFrequency,
                                             &SionnaRtChannelModel::GetFrequency),
                          MakeDoubleChecker<double>())
            .AddAttribute("UpdatePeriod",
                          "Specify the channel coherence time",
                          TimeValue(MilliSeconds(0)),
                          MakeTimeAccessor(&SionnaRtChannelModel::m_updatePeriod),
                          MakeTimeChecker())
            .AddAttribute("OutputImageDirectory",
                          "Output directory for rendered Sionna RT images",
                          StringValue("images"),
                          MakeStringAccessor(&SionnaRtChannelModel::m_outputImageDirectory),
                          MakeStringChecker())
            .AddAttribute("OutputImageName",
                          "Base file name for rendered Sionna RT images",
                          StringValue("screenshot"),
                          MakeStringAccessor(&SionnaRtChannelModel::m_outputImageName),
                          MakeStringChecker())
            .AddAttribute("IsImageRenderedEnabled",
                          "Enable or disable scene image rendering",
                          BooleanValue(false),
                          MakeBooleanAccessor(&SionnaRtChannelModel::m_isImageRendered),
                          MakeBooleanChecker())
            .AddAttribute("CameraPosition",
                          "Camera position for scene rendering",
                          VectorValue(Vector(0.0, 0.0, 300.0)),
                          MakeVectorAccessor(&SionnaRtChannelModel::m_cameraPosition),
                          MakeVectorChecker())
            .AddAttribute("CameraLookAt",
                          "Camera look-at position for scene rendering",
                          VectorValue(Vector(0.0, 0.0, 0.0)),
                          MakeVectorAccessor(&SionnaRtChannelModel::m_cameraLookAt),
                          MakeVectorChecker())
            .AddAttribute("IsMergeShapeEnabled",
                          "Enable or disable MEGE shape loading",
                          BooleanValue(false),
                          MakeBooleanAccessor(&SionnaRtChannelModel::SetMergeShapeEnable),
                          MakeBooleanChecker())
            .AddAttribute(
                "MaxNumberOfPaths",
                "The maximum number of paths to be generated by the Sionna RT path solver",
                DoubleValue(90),
                MakeDoubleAccessor(&SionnaRtChannelModel::SetMaxNumberOfPaths,
                                   &SionnaRtChannelModel::GetMaxNumberOfPaths),
                MakeDoubleChecker<double>())
            .AddAttribute("NormalizeDelays",
                          "Enable normalization of path delays in Sionna RT CIR output",
                          BooleanValue(true),
                          MakeBooleanAccessor(&SionnaRtChannelModel::SetNormalizeDelays,
                                              &SionnaRtChannelModel::GetNormalizeDelays),
                          MakeBooleanChecker())
            .AddAttribute("Scenario",
                          "The Sionna RT scenario name used to create scenes",
                          StringValue("munich"),
                          MakeStringAccessor(&SionnaRtChannelModel::SetScenario,
                                             &SionnaRtChannelModel::GetScenario),
                          MakeStringChecker());
    return tid;
}

void
SionnaRtChannelModel::SetMergeShapeEnable(bool cond)
{
    NS_LOG_FUNCTION(this);
    m_sceneConfigs.mergeShapes = cond;
}

void
SionnaRtChannelModel::SetFrequency(double f)
{
    NS_LOG_FUNCTION(this);
    NS_ASSERT_MSG(f >= 500.0e6 && f <= 100.0e9,
                  "Frequency should be between 0.5 and 100 GHz but is " << f);
    NS_LOG_INFO("Setting operating frequency to " << f << " Hz");
    m_sceneConfigs.frequency = f;
}

void
SionnaRtChannelModel::SetMaxNumberOfPaths(double p)
{
    NS_LOG_FUNCTION(this);
    NS_LOG_INFO("Setting max number of paths to " << p);
    m_maxNumberOfPaths = p;
}

void
SionnaRtChannelModel::SetNormalizeDelays(bool cond)
{
    NS_LOG_FUNCTION(this);
    NS_LOG_INFO("Setting normalize delays to " << (cond ? "true" : "false"));
    m_normalizeDelays = cond;
}

bool
SionnaRtChannelModel::GetNormalizeDelays() const
{
    NS_LOG_FUNCTION(this);
    return m_normalizeDelays;
}

double
SionnaRtChannelModel::GetMaxNumberOfPaths() const
{
    NS_LOG_FUNCTION(this);
    NS_LOG_INFO("Getting max number of paths: " << m_maxNumberOfPaths);
    return m_maxNumberOfPaths;
}

double
SionnaRtChannelModel::GetFrequency() const
{
    NS_LOG_FUNCTION(this);
    return m_sceneConfigs.frequency;
}

bool
SionnaRtChannelModel::AntennaSetupChanged(Ptr<const PhasedArrayModel> aAntenna,
                                          Ptr<const PhasedArrayModel> bAntenna,
                                          Ptr<const ChannelMatrix> channelMatrix) const
{
    // This allows changing the number of antenna ports during execution,
    // which is used by nr's initial association.
    size_t sAntNumElems = aAntenna->GetNumElems();
    size_t uAntNumElems = bAntenna->GetNumElems();
    size_t chanNumRows = channelMatrix->m_channel.GetNumRows();
    size_t chanNumCols = channelMatrix->m_channel.GetNumCols();

    // Log a description of what changed if any
    if (((uAntNumElems != chanNumRows) || (sAntNumElems != chanNumCols)) &&
        ((uAntNumElems != chanNumCols) || (sAntNumElems != chanNumRows)))
    {
        NS_LOG_INFO("Antenna setup changed: aAntenna elems="
                    << sAntNumElems << " bAntenna elems=" << uAntNumElems << " channel dims=("
                    << chanNumRows << ", " << chanNumCols << ")");
    }
    bool changed = (((uAntNumElems != chanNumRows) || (sAntNumElems != chanNumCols)) &&
                    ((uAntNumElems != chanNumCols) || (sAntNumElems != chanNumRows))) ||
                   aAntenna->IsChannelOutOfDate(bAntenna);

    if (changed)
    {
        NS_LOG_DEBUG("AntennaSetupChanged evaluated to true");
    }
    return changed;
}

bool
SionnaRtChannelModel::ChannelMatrixNeedsUpdate(Ptr<const ChannelMatrix> channelMatrix) const
{
    NS_LOG_FUNCTION(this);

    bool update = false;

    // if the coherence time is over the channel has to be updated
    if (!m_updatePeriod.IsZero() &&
        Simulator::Now() - channelMatrix->m_generatedTime > m_updatePeriod)
    {
        NS_LOG_DEBUG("Generation time " << channelMatrix->m_generatedTime.As(Time::NS) << " now "
                                        << Now().As(Time::NS));
        update = true;
    }

    return update;
}

Ptr<const MatrixBasedChannelModel::ChannelMatrix>
SionnaRtChannelModel::GetChannel(Ptr<const MobilityModel> aMob,
                                 Ptr<const MobilityModel> bMob,
                                 Ptr<const PhasedArrayModel> aAntenna,
                                 Ptr<const PhasedArrayModel> bAntenna)
{
    NS_LOG_FUNCTION(this);

    // Compute the channel matrix key. The key is reciprocal, i.e., key (a, b) = key (b, a)
    uint64_t channelMatrixKey = GetKey(aAntenna->GetId(), bAntenna->GetId());

    bool notFound = false;
    bool update = false;
    Ptr<ChannelMatrix> channelMatrix;

    if (m_channelMatrixMap.find(channelMatrixKey) != m_channelMatrixMap.end())
    {
        // channel matrix present in the map
        NS_LOG_DEBUG("channel matrix present in the map");
        channelMatrix = m_channelMatrixMap[channelMatrixKey];
        update = ChannelMatrixNeedsUpdate(channelMatrix);
        update |= AntennaSetupChanged(aAntenna, bAntenna, channelMatrix);
        if (update)
        {
            NS_LOG_INFO("Cached channel matrix marked for update (stale or antenna changed)");
        }
    }
    else
    {
        NS_LOG_DEBUG("channel matrix not found");
        notFound = true;
    }

    // If the channel is not present in the map or if it has to be updated
    // generate a new realization
    if (notFound || update)
    {
        // channel matrix not found or has to be updated, generate a new one
        py::module_ rt = py::module_::import("sionna.rt");
        NS_LOG_INFO("Building scene for antenna pair ("
                    << aAntenna->GetId() << ", " << bAntenna->GetId() << ") and nodes ("
                    << aMob->GetObject<Node>()->GetId() << ", " << bMob->GetObject<Node>()->GetId()
                    << ")");

        py::object scene = CreateScene(rt, aMob, bMob, aAntenna, bAntenna);
        py::object paths = CalculatePaths(rt, scene);

        channelMatrix = GetNewChannel(paths, aMob, bMob, aAntenna, bAntenna);
        channelMatrix->m_antennaPair =
            std::make_pair(aAntenna->GetId(),
                           bAntenna->GetId()); // save antenna pair, with the exact order of s and u
                                               // antennas at the moment of the channel generation

        // store or replace the channel matrix in the channel map
        m_channelMatrixMap[channelMatrixKey] = channelMatrix;

        // each time a new channel is created, also compute the channel params
        auto channelParams = CalculateChannelParamsFromPaths(paths, aMob, bMob);
        uint64_t channelParamsKey =
            GetKey(aMob->GetObject<Node>()->GetId(), bMob->GetObject<Node>()->GetId());
        m_channelParamsMap[channelParamsKey] = channelParams;
    }

    return channelMatrix;
}

Ptr<MatrixBasedChannelModel::ChannelMatrix>
SionnaRtChannelModel::GetNewChannel(const py::object paths,
                                    const Ptr<const MobilityModel> sMob,
                                    const Ptr<const MobilityModel> uMob,
                                    Ptr<const PhasedArrayModel> sAntenna,
                                    Ptr<const PhasedArrayModel> uAntenna) const
{
    NS_LOG_FUNCTION(this);
    NS_LOG_DEBUG("Creating new ChannelMatrix from path information");
    Complex3DVector hUsn = CalculateCirFromPaths(paths);

    Ptr<ChannelMatrix> channelMatrix = CreateNewMatrixChannel(hUsn, sMob, uMob, sAntenna, uAntenna);
    NS_LOG_DEBUG("Successfully created new ChannelMatrix");
    return channelMatrix;
}

Ptr<MatrixBasedChannelModel::ChannelMatrix>
SionnaRtChannelModel::CreateNewMatrixChannel(const Complex3DVector hUsn,
                                             const Ptr<const MobilityModel> sMob,
                                             const Ptr<const MobilityModel> uMob,
                                             Ptr<const PhasedArrayModel> sAntenna,
                                             Ptr<const PhasedArrayModel> uAntenna) const
{
    // create a channel matrix instance
    Ptr<ChannelMatrix> channelMatrix = Create<ChannelMatrix>();
    channelMatrix->m_generatedTime = Simulator::Now();
    // save in which order is generated this matrix
    channelMatrix->m_nodeIds =
        std::make_pair(sMob->GetObject<Node>()->GetId(), uMob->GetObject<Node>()->GetId());

    channelMatrix->m_channel = hUsn;
    return channelMatrix;
}

void
SionnaRtChannelModel::SetScenario(const std::string& scenario)
{
    NS_LOG_FUNCTION(this);

    // Check if scenario is an XML file path (ends with .xml)
    bool isXmlFile = (scenario.length() >= 4 && scenario.substr(scenario.length() - 4) == ".xml");

    if (!isXmlFile && builtInSceneSionna.count(scenario) == 0)
    {
        NS_FATAL_ERROR("Unsupported Sionna scenario name: '"
                       << scenario << "'. Allowed scenarios are: "
                       << "box, box_one_screen, box_two_screens, double_reflector, "
                       << "triple_reflector, simple_wedge, Etoile, floor_wall, Florence, "
                       << "Munich, san_francisco, simple_street_canyon, "
                       << "simple_street_canyon_with_cars, simple_reflector, "
                       << "or a custom XML file path ending with .xml");
    }

    NS_LOG_INFO("Changing scenario to: " << scenario
                                         << (isXmlFile ? " (XML file)" : " (built-in)"));

    m_sceneConfigs.sceneName = scenario;
}

std::string
SionnaRtChannelModel::GetScenario() const
{
    NS_LOG_FUNCTION(this);
    return m_sceneConfigs.sceneName;
}

py::object
SionnaRtChannelModel::LoadScene(const py::module_ rt) const
{
    NS_LOG_FUNCTION(this);
    NS_LOG_INFO("Loading Sionna RT scene: " << m_sceneConfigs.sceneName << " mergeShapes="
                                            << (m_sceneConfigs.mergeShapes ? "true" : "false"));

    py::object scene;

    // Check if scenario is an XML file path (ends with .xml)
    bool isXmlFile =
        (m_sceneConfigs.sceneName.length() >= 4 &&
         m_sceneConfigs.sceneName.substr(m_sceneConfigs.sceneName.length() - 4) == ".xml");

    if (isXmlFile)
    {
        // Load from XML file using Scene constructor
        NS_LOG_DEBUG("Loading scene from XML file: " << m_sceneConfigs.sceneName);
        py::object Scene = rt.attr("Scene");
        scene = Scene(py::arg("filename") = m_sceneConfigs.sceneName,
                      py::arg("merge_shapes") = m_sceneConfigs.mergeShapes);
    }
    else
    {
        // Load built-in scenario using load_scene
        NS_LOG_DEBUG("Loading built-in Sionna RT scenario: " << m_sceneConfigs.sceneName);
        py::object load_scene = rt.attr("load_scene");
        py::object scenes = rt.attr("scene");
        py::object scene_name = scenes.attr(m_sceneConfigs.sceneName.c_str());
        scene = load_scene(scene_name, py::arg("merge_shapes") = m_sceneConfigs.mergeShapes);
    }

    NS_LOG_DEBUG("Scene object loaded from Sionna RT");
    return scene;
}

std::string
SionnaRtChannelModel::GetAntennaElementPattern(Ptr<const AntennaModel> element) const
{
    NS_LOG_FUNCTION(this);
    std::string pattern;

    if (DynamicCast<const ThreeGppAntennaModel>(element) != nullptr)
    {
        pattern = "tr38901";
    }
    else
    {
        pattern = "iso";
    }
    NS_LOG_DEBUG("Antenna element pattern: " << pattern);
    return pattern;
}

std::string
SionnaRtChannelModel::GetPolarizationFromPolSlantAngle(const double PolSlantAngle,
                                                       const uint8_t numpol) const
{
    NS_LOG_FUNCTION(this);
    std::string polarization = "V";
    if (numpol == 1)
    {
        if (PolSlantAngle == 0.0)
        {
            polarization = "V";
        }
        else if (PolSlantAngle == M_PI / 2)
        {
            polarization = "H";
        }
    }
    else
    {
        if (PolSlantAngle == 0.0)
        {
            polarization = "VH";
        }
        else
        {
            polarization = "cross";
        }
    }
    return polarization;
}

py::object
SionnaRtChannelModel::MakePlannarArray(const py::module_ rt,
                                       Ptr<const PhasedArrayModel> antenna) const
{
    NS_LOG_FUNCTION(this);
    py::object PlanarArray = rt.attr("PlanarArray");
    auto upa = DynamicCast<const UniformPlanarArray>(antenna);
    if (!upa)
    {
        NS_LOG_WARN("Antenna is not UniformPlanarArray; using defaults 1x1,tr38901, V-pol.");
        return PlanarArray(py::arg("num_rows") = 1,
                           py::arg("num_cols") = 1,
                           py::arg("pattern") = "tr38901",
                           py::arg("polarization") = "V");
    }

    return PlanarArray(py::arg("num_rows") = upa->GetNumRows(),
                       py::arg("num_cols") = upa->GetNumColumns(),
                       py::arg("vertical_spacing") = upa->GetAntennaVerticalSpacing(),
                       py::arg("horizontal_spacing") = upa->GetAntennaHorizontalSpacing(),
                       py::arg("pattern") = GetAntennaElementPattern(upa->GetAntennaElement()),
                       py::arg("polarization") =
                           GetPolarizationFromPolSlantAngle(upa->GetPolSlant(), upa->GetNumPols()));
}

Ptr<const MatrixBasedChannelModel::ChannelParams>
SionnaRtChannelModel::GetParams(Ptr<const MobilityModel> aMob, Ptr<const MobilityModel> bMob) const
{
    NS_LOG_FUNCTION(this);

    // Compute the channel key. The key is reciprocal, i.e., key (a, b) = key (b, a)
    uint64_t channelParamsKey =
        GetKey(aMob->GetObject<Node>()->GetId(), bMob->GetObject<Node>()->GetId());

    if (m_channelParamsMap.find(channelParamsKey) != m_channelParamsMap.end())
    {
        return m_channelParamsMap.find(channelParamsKey)->second;
    }

    NS_LOG_WARN("Channel params map not found. Returning a nullptr.");
    return nullptr;
}

Ptr<SionnaRtChannelModel::SionnaRtChannelParams>
SionnaRtChannelModel::CalculateChannelParamsFromPaths(const py::object paths,
                                                      Ptr<const MobilityModel> aMob,
                                                      Ptr<const MobilityModel> bMob) const
{
    NS_LOG_FUNCTION(this);
    NS_LOG_DEBUG("Calculating ChannelParams from paths");
    Ptr<SionnaRtChannelModel::SionnaRtChannelParams> channelParams =
        Create<SionnaRtChannelModel::SionnaRtChannelParams>();

    channelParams->m_nodeIds =
        std::make_pair(aMob->GetObject<Node>()->GetId(), bMob->GetObject<Node>()->GetId());

    py::module_ np = py::module_::import("numpy");

    channelParams->m_generatedTime = Simulator::Now();

    channelParams->m_delay = CalculateTauFromPaths(np, paths);
    channelParams->m_doppler = CalculateDopplerFromPaths(np, paths);
    channelParams->m_angle.clear();
    channelParams->m_angle.push_back(CalculateAnglesFromPaths(np, paths, "AOA")); // AOA
    channelParams->m_angle.push_back(CalculateAnglesFromPaths(np, paths, "ZOA")); // ZOA
    channelParams->m_angle.push_back(CalculateAnglesFromPaths(np, paths, "AOD")); // AOD
    channelParams->m_angle.push_back(CalculateAnglesFromPaths(np, paths, "ZOD")); // ZOD

    // Precompute angles sincos
    channelParams->m_cachedAngleSincos.resize(channelParams->m_angle.size());
    for (size_t direction = 0; direction < channelParams->m_angle.size(); direction++)
    {
        channelParams->m_cachedAngleSincos[direction].resize(
            channelParams->m_angle[direction].size());
        for (size_t cluster = 0; cluster < channelParams->m_angle[direction].size(); cluster++)
        {
            channelParams->m_cachedAngleSincos[direction][cluster] = {
                sin(channelParams->m_angle[direction][cluster] * DEG2RAD),
                cos(channelParams->m_angle[direction][cluster] * DEG2RAD)};
        }
    }

    return channelParams;
}

py::object
SionnaRtChannelModel::CreateScene(const py::module_ rt,
                                  const Ptr<const MobilityModel> sMob,
                                  const Ptr<const MobilityModel> uMob,
                                  Ptr<const PhasedArrayModel> sAntenna,
                                  Ptr<const PhasedArrayModel> uAntenna) const
{
    NS_LOG_FUNCTION(this);
    py::object scene = LoadScene(rt);
    NS_LOG_DEBUG("Installing antenna arrays on scene (tx/rx)");
    scene.attr("tx_array") = MakePlannarArray(rt, sAntenna);
    scene.attr("rx_array") = MakePlannarArray(rt, uAntenna);
    // --- Create and add a TX/RX

    py::object Transmitter = rt.attr("Transmitter");
    py::object Receiver = rt.attr("Receiver");
    Vector pos = sMob->GetPosition();
    Vector v = sMob->GetVelocity();
    auto upa = DynamicCast<const UniformPlanarArray>(sAntenna);
    Vector posRx = uMob->GetPosition();
    py::object tx =
        Transmitter(py::arg("name") = "tx",
                    py::arg("position") = py::make_tuple(pos.x, pos.y, pos.z),
                    py::arg("orientation") = py::make_tuple(upa->GetAlpha(), upa->GetBeta(), 0.0),
                    py::arg("velocity") = py::make_tuple(v.x, v.y, v.z));

    Vector vRx = uMob->GetVelocity();
    auto upau = DynamicCast<const UniformPlanarArray>(uAntenna);
    py::object rx =
        Receiver(py::arg("name") = "rx",
                 py::arg("position") = py::make_tuple(posRx.x, posRx.y, posRx.z),
                 py::arg("orientation") = py::make_tuple(upau->GetAlpha(), upau->GetBeta(), 0.0),
                 py::arg("velocity") = py::make_tuple(vRx.x, vRx.y, vRx.z));

    scene.attr("add")(tx);
    scene.attr("add")(rx);

    //  tx.attr("look_at")(rx);
    scene.attr("frequency") = m_sceneConfigs.frequency;
    NS_LOG_INFO("Scene created: frequency=" << m_sceneConfigs.frequency << " tx@" << pos << " rx@"
                                            << posRx);

    return scene;
}

void
SionnaRtChannelModel::SetRtPathSolverConfig(SionnaRtChannelModel::RtPathSolverConfig configs)
{
    m_RtPathSolverConfig = configs;
    NS_LOG_INFO("RtPathSolverConfig set: maxDepth="
                << configs.maxDepth << " los=" << configs.los << " specular_reflection="
                << configs.specularReflection << " diffuse_reflection=" << configs.diffuseReflection
                << " refraction=" << configs.refraction << " diffraction=" << configs.diffraction
                << " edgeDiffraction=" << configs.edgeDiffraction
                << " synthetic_array=" << configs.syntheticArray << " seed=" << configs.seed);
}

SionnaRtChannelModel::RtPathSolverConfig
SionnaRtChannelModel::GetRtPathSolverConfig() const
{
    return m_RtPathSolverConfig;
}

py::object
SionnaRtChannelModel::CalculatePaths(const py::module_ rt, const py::object scene)
{
    NS_LOG_FUNCTION(this);
    py::object PathSolver = rt.attr("PathSolver");
    py::object psolver = PathSolver();

    static std::string rt_version = py::str(rt.attr("__version__"));
    NS_LOG_INFO("Using Sionna RT version: " << rt_version);

    py::object paths;
    NS_LOG_DEBUG("Invoking Sionna RT PathSolver with configured options");

    if (rt_version == "1.1.0")
    {
        NS_LOG_DEBUG("Using Sionna RT PathSolver API == 1.1.0");
        paths = psolver(py::arg("scene") = scene,
                        py::arg("max_depth") = m_RtPathSolverConfig.maxDepth,
                        py::arg("los") = m_RtPathSolverConfig.los,
                        py::arg("specular_reflection") = m_RtPathSolverConfig.specularReflection,
                        py::arg("diffuse_reflection") = m_RtPathSolverConfig.diffuseReflection,
                        py::arg("refraction") = m_RtPathSolverConfig.refraction,
                        py::arg("synthetic_array") = m_RtPathSolverConfig.syntheticArray,
                        py::arg("seed") = m_RtPathSolverConfig.seed);
    }
    else
    {
        NS_LOG_DEBUG("Using Sionna RT PathSolver API >= 1.2.0 used");
        paths = psolver(py::arg("scene") = scene,
                        py::arg("max_depth") = m_RtPathSolverConfig.maxDepth,
                        py::arg("los") = m_RtPathSolverConfig.los,
                        py::arg("specular_reflection") = m_RtPathSolverConfig.specularReflection,
                        py::arg("diffuse_reflection") = m_RtPathSolverConfig.diffuseReflection,
                        py::arg("refraction") = m_RtPathSolverConfig.refraction,
                        py::arg("diffraction") = m_RtPathSolverConfig.diffraction,
                        py::arg("edge_diffraction") = m_RtPathSolverConfig.edgeDiffraction,
                        py::arg("synthetic_array") = m_RtPathSolverConfig.syntheticArray,
                        py::arg("seed") = m_RtPathSolverConfig.seed);
    }

    NS_LOG_INFO("Path computation finished");

    size_t num_paths = GetNumberOfPathsFromCir(paths);
    NS_LOG_INFO("Number of generated paths: " << num_paths);
    if (DoesNumberOfPathsExceedMaximum(num_paths))
    {
        NS_LOG_WARN("Number of generated paths (" << num_paths
                                                  << ") exceeds the configured maximum number of "
                                                     "paths ("
                                                  << m_maxNumberOfPaths
                                                  << "). The excess paths will be ignored.");

        m_RtPathSolverConfig.maxDepth -= 1;
        paths = CalculatePaths(rt, scene);
    }

    if (m_isImageRendered && num_paths > 0)
    {
        NS_LOG_DEBUG("Rendering scene image to disk");
        SceneRenderImageToFile(rt, paths, scene);
    }
    return paths;
}

void
SionnaRtChannelModel::SceneRenderImageToFile(const py::module_ rt,
                                             const py::object paths,
                                             const py::object scene) const
{
    static int count = 0;
    if (!std::filesystem::exists(m_outputImageDirectory + "/"))
    {
        std::filesystem::create_directories(m_outputImageDirectory + "/");
    }

    py::object camera = rt.attr("Camera");
    py::object Mycamera = camera(
        py::arg("position") =
            py::make_tuple(m_cameraPosition.x, m_cameraPosition.y, m_cameraPosition.z),
        py::arg("look_at") = py::make_tuple(m_cameraLookAt.x, m_cameraLookAt.y, m_cameraLookAt.z));
    py::object renderer = scene.attr("render_to_file");
    py::object image =
        renderer(py::arg("camera") = Mycamera,
                 py::arg("filename") = m_outputImageDirectory + "/" + m_outputImageName +
                                       std::to_string(count) + ".png",
                 py::arg("paths") = paths);
    NS_LOG_INFO("Scene image rendered");
    count += 1;
}

MatrixBasedChannelModel::Complex3DVector
SionnaRtChannelModel::CalculateCirFromPaths(const py::object paths) const
{
    NS_LOG_FUNCTION(this);
    py::module_ np = py::module_::import("numpy");

    py::tuple cir = paths.attr("cir")(py::arg("normalize_delays") = m_normalizeDelays,
                                      py::arg("out_type") = "numpy");

    py::object a_py = np.attr("asarray")(cir[0], py::arg("dtype") = np.attr("complex128"));
    py::buffer_info a_info = py::buffer(a_py).request();
    auto* data = static_cast<std::complex<double>*>(a_info.ptr);

    // Generate channel coefficients for each cluster n and each receiver
    //  and transmitter element pair u,s.
    // where n is cluster index, u and s are receive and transmit antenna element.
    // Create the Complex3DVector and copy data a_info.shape[1]:num_ant_rx,
    // a_info.shape[3]:num_ant_tx, a_info.shape[4]:num_paths
    Complex3DVector hUsn(a_info.shape[1], a_info.shape[3], a_info.shape[4]);
    NS_LOG_DEBUG("CIR numpy buffer shapes: shape[1]="
                 << a_info.shape[1] << " shape[3]=" << a_info.shape[3]
                 << " shape[4]=" << a_info.shape[4] << " total size=" << a_info.size);

    for (size_t cIndex = 0; cIndex < hUsn.GetNumPages(); cIndex++)
    {
        for (size_t rowIdx = 0; rowIdx < hUsn.GetNumRows(); rowIdx++)
        {
            for (size_t colIdx = 0; colIdx < hUsn.GetNumCols(); colIdx++)
            {
                size_t idx = (cIndex * hUsn.GetNumRows() + rowIdx) * hUsn.GetNumCols() + colIdx;
                hUsn(rowIdx, colIdx, cIndex) = data[idx];
            }
        }
    }

    return hUsn;
}

MatrixBasedChannelModel::DoubleVector
SionnaRtChannelModel::CalculateTauFromPaths(const py::module_ np, py::object paths) const
{
    NS_LOG_FUNCTION(this);
    // Convert to numpy float64 array
    py::tuple cir = paths.attr("cir")(py::arg("normalize_delays") = m_normalizeDelays,
                                      py::arg("out_type") = "numpy");
    py::object tau_py = np.attr("asarray")(cir[1], py::arg("dtype") = np.attr("float64"));
    // Buffer info
    py::buffer_info tau_info = py::buffer(tau_py).request();
    auto* data = static_cast<double*>(tau_info.ptr);
    // Shape of array in case syntheticarray differs in sionna
    if (m_RtPathSolverConfig.syntheticArray)
    {
        return DoubleVector(data, data + tau_info.shape[2]);
    }
    else
    {
        return DoubleVector(data, data + tau_info.shape[4]);
    }
}

std::vector<double>
SionnaRtChannelModel::CalculateDopplerFromPaths(const py::module_ np, py::object paths) const
{
    NS_LOG_FUNCTION(this);
    py::tuple doppler = paths.attr("doppler");
    // Convert doppler (DrJit tensor) to numpy complex128 array
    py::object d_py = np.attr("asarray")(doppler, py::arg("dtype") = np.attr("float64"));

    // Buffer info
    py::buffer_info d_info = py::buffer(d_py).request();
    auto* data = static_cast<double*>(d_info.ptr);

    // Shape of array in case syntheticarray differs in sionna
    size_t dopplerSize =
        (m_RtPathSolverConfig.syntheticArray) ? size_t(d_info.shape[2]) : size_t(d_info.shape[4]);

    std::vector<double> dopplerVector(dopplerSize);
    for (size_t i = 0; i < dopplerSize; i++)
    {
        dopplerVector[i] = data[i];
    }

    return dopplerVector;
}

MatrixBasedChannelModel::DoubleVector
SionnaRtChannelModel::CalculateAnglesFromPaths(const py::module_ np,
                                               py::object paths,
                                               std::string Angle) const
{
    py::object angle_py;
    if (Angle == "ZOA")
    {
        angle_py = paths.attr("theta_r");
    }
    else if (Angle == "ZOD")
    {
        angle_py = paths.attr("theta_t");
    }
    else if (Angle == "AOA")
    {
        angle_py = paths.attr("phi_r");
    }
    else if (Angle == "AOD")
    {
        angle_py = paths.attr("phi_t");
    }

    else
    {
        NS_LOG_ERROR("Angle type not recognized. Use ZOA, or ZOD, or AOA, or AOD.");
    }

    // Convert to numpy float64 array
    py::object angle_np = np.attr("asarray")(angle_py, py::arg("dtype") = py::dtype("float64"));
    // Buffer info

    py::buffer_info angle_info = py::buffer(angle_np).request();
    auto* data = static_cast<double*>(angle_info.ptr);

    // Shape of array in case syntheticarray differs in sionna
    size_t angelSize = (m_RtPathSolverConfig.syntheticArray) ? size_t(angle_info.shape[2])
                                                             : size_t(angle_info.shape[4]);

    MatrixBasedChannelModel::DoubleVector angleVector;
    for (size_t i = 0; i < angelSize; ++i)
    {
        angleVector.push_back(data[i]);
    }
    NS_LOG_DEBUG("Extracted " << angleVector.size() << " values for angle type " << Angle);

    return angleVector;
}

size_t
SionnaRtChannelModel::GetNumberOfPathsFromCir(const py::object paths) const
{
    NS_LOG_FUNCTION(this);
    py::module_ np = py::module_::import("numpy");

    py::object angle_np =
        np.attr("asarray")(paths.attr("phi_r"), py::arg("dtype") = py::dtype("float64"));
    py::buffer_info angle_info = py::buffer(angle_np).request();

    // Get number of paths from shape[4] (or shape[2] if synthetic array)
    size_t num_paths =
        m_RtPathSolverConfig.syntheticArray ? angle_info.shape[2] : angle_info.shape[4];
    return num_paths;
}

bool
SionnaRtChannelModel::DoesNumberOfPathsExceedMaximum(size_t numPaths) const
{
    NS_LOG_FUNCTION(this << numPaths);
    return (numPaths > static_cast<size_t>(m_maxNumberOfPaths));
}

} // namespace ns3
