#include "OpenSimLoader.hpp"

#include <OpenSim/Common/Storage.h>
#include <OpenSim/Simulation/Model/Appearance.h>
#include <OpenSim/Simulation/Model/Geometry.h>
#include <OpenSim/Simulation/Model/Marker.h>
#include <OpenSim/Simulation/Model/Model.h>
#include <OpenSim/Simulation/Model/ModelVisualizer.h>
#include <OpenSim/Simulation/StatesTrajectory.h>
#include <SimTKcommon/internal/PolygonalMesh.h>

#include <Magnum/Math/Functions.h>
#include <Magnum/Math/Matrix4.h>

#include <algorithm>
#include <cmath>
#include <filesystem>
#include <iostream>
#include <limits>
#include <memory>
#include <set>
#include <stdexcept>
#include <string>
#include <unordered_map>

using namespace Magnum;

namespace osim_viewer {

namespace {

Vector3 toVector3(const SimTK::Vec3& value) {
    return {Float(value[0]), Float(value[1]), Float(value[2])};
}

Quaternion toQuaternion(const SimTK::Rotation& rotation) {
    const SimTK::Quaternion quat = rotation.convertRotationToQuaternion();
    return Quaternion{{Float(quat[1]), Float(quat[2]), Float(quat[3])}, Float(quat[0])}.normalized();
}

Color4 toColor(const OpenSim::Appearance& appearance) {
    const SimTK::Vec3 color = appearance.get_color();
    return {Float(color[0]), Float(color[1]), Float(color[2]), Float(appearance.get_opacity())};
}

Range3D scaledBounds(const Range3D& bounds, const Vector3& scale) {
    Range3D out;
    out.min() = Vector3{std::numeric_limits<Float>::max()};
    out.max() = Vector3{-std::numeric_limits<Float>::max()};
    for(const Vector3& corner: {
            Vector3{bounds.min().x(), bounds.min().y(), bounds.min().z()},
            Vector3{bounds.min().x(), bounds.min().y(), bounds.max().z()},
            Vector3{bounds.min().x(), bounds.max().y(), bounds.min().z()},
            Vector3{bounds.min().x(), bounds.max().y(), bounds.max().z()},
            Vector3{bounds.max().x(), bounds.min().y(), bounds.min().z()},
            Vector3{bounds.max().x(), bounds.min().y(), bounds.max().z()},
            Vector3{bounds.max().x(), bounds.max().y(), bounds.min().z()},
            Vector3{bounds.max().x(), bounds.max().y(), bounds.max().z()},
        }) {
        out = Math::join(out, Range3D::fromSize(corner*scale, {}));
    }
    return out;
}

Range3D transformBounds(const Range3D& bounds, const Matrix4& transform) {
    Range3D out;
    out.min() = Vector3{std::numeric_limits<Float>::max()};
    out.max() = Vector3{-std::numeric_limits<Float>::max()};
    for(const Vector3& corner: {
            Vector3{bounds.min().x(), bounds.min().y(), bounds.min().z()},
            Vector3{bounds.min().x(), bounds.min().y(), bounds.max().z()},
            Vector3{bounds.min().x(), bounds.max().y(), bounds.min().z()},
            Vector3{bounds.min().x(), bounds.max().y(), bounds.max().z()},
            Vector3{bounds.max().x(), bounds.min().y(), bounds.min().z()},
            Vector3{bounds.max().x(), bounds.min().y(), bounds.max().z()},
            Vector3{bounds.max().x(), bounds.max().y(), bounds.min().z()},
            Vector3{bounds.max().x(), bounds.max().y(), bounds.max().z()},
        }) {
        out = Math::join(out, Range3D::fromSize(transform.transformPoint(corner), {}));
    }
    return out;
}

struct MeshCache {
    std::unordered_map<std::string, std::shared_ptr<CpuMesh>> meshes;
    std::shared_ptr<CpuMesh> unitSphere;
    std::shared_ptr<CpuMesh> unitCylinder;
    std::shared_ptr<CpuMesh> unitBrick;
};

std::shared_ptr<CpuMesh> buildCpuMesh(const SimTK::PolygonalMesh& mesh) {
    auto cpu = std::make_shared<CpuMesh>();
    std::vector<Vector3> positions;
    positions.reserve(mesh.getNumVertices());
    for(int i = 0; i != mesh.getNumVertices(); ++i) positions.push_back(toVector3(mesh.getVertexPosition(i)));

    std::vector<Vector3> normals(positions.size(), Vector3{0.0f});
    for(int face = 0; face != mesh.getNumFaces(); ++face) {
        const int faceVertexCount = mesh.getNumVerticesForFace(face);
        if(faceVertexCount < 3) continue;
        const int a = mesh.getFaceVertex(face, 0);
        for(int i = 1; i < faceVertexCount - 1; ++i) {
            const int b = mesh.getFaceVertex(face, i);
            const int c = mesh.getFaceVertex(face, i + 1);
            const Vector3 ab = positions[b] - positions[a];
            const Vector3 ac = positions[c] - positions[a];
            const Vector3 faceNormal = Math::cross(ab, ac);
            normals[a] += faceNormal;
            normals[b] += faceNormal;
            normals[c] += faceNormal;
            cpu->indices.push_back(Magnum::UnsignedInt(a));
            cpu->indices.push_back(Magnum::UnsignedInt(b));
            cpu->indices.push_back(Magnum::UnsignedInt(c));
        }
    }

    Range3D bounds;
    bounds.min() = Vector3{std::numeric_limits<Float>::max()};
    bounds.max() = Vector3{-std::numeric_limits<Float>::max()};

    cpu->vertices.reserve(positions.size());
    for(std::size_t i = 0; i != positions.size(); ++i) {
        const Vector3 normal = normals[i].isZero() ? Vector3::yAxis() : normals[i].normalized();
        cpu->vertices.push_back({positions[i], normal});
        bounds = Math::join(bounds, Range3D::fromSize(positions[i], {}));
    }
    cpu->bounds = bounds;
    return cpu;
}

std::shared_ptr<CpuMesh> sphereMesh(MeshCache& cache) {
    if(!cache.unitSphere) {
        cache.unitSphere = buildCpuMesh(SimTK::PolygonalMesh::createSphereMesh(1.0, 2));
    }
    return cache.unitSphere;
}

std::shared_ptr<CpuMesh> cylinderMesh(MeshCache& cache) {
    if(!cache.unitCylinder) {
        cache.unitCylinder = buildCpuMesh(SimTK::PolygonalMesh::createCylinderMesh(SimTK::YAxis, 1.0, 1.0, 2));
    }
    return cache.unitCylinder;
}

std::shared_ptr<CpuMesh> brickMesh(MeshCache& cache) {
    if(!cache.unitBrick) {
        cache.unitBrick = buildCpuMesh(SimTK::PolygonalMesh::createBrickMesh(SimTK::Vec3(1.0, 1.0, 1.0), 1));
    }
    return cache.unitBrick;
}

std::shared_ptr<CpuMesh> resolveMeshFile(MeshCache& cache, const OpenSim::Model& model, const std::string& meshFile) {
    bool isAbsolute = false;
    SimTK::Array_<std::string> attempts;
    if(!OpenSim::ModelVisualizer::findGeometryFile(model, meshFile, isAbsolute, attempts) || attempts.empty()) {
        throw std::runtime_error("Unable to locate geometry file: " + meshFile);
    }

    const std::string resolved = attempts.back();
    if(auto found = cache.meshes.find(resolved); found != cache.meshes.end()) return found->second;

    SimTK::PolygonalMesh mesh;
    mesh.loadFile(resolved);
    auto cpu = buildCpuMesh(mesh);
    cache.meshes.emplace(resolved, cpu);
    return cpu;
}

Vector3 geometryScale(const OpenSim::Geometry& geometry) {
    const SimTK::Vec3 scale = geometry.get_scale_factors();
    return {Float(scale[0]), Float(scale[1]), Float(scale[2])};
}

GeometrySpec makeGeometrySpec(MeshCache& cache, const OpenSim::Model& model, const OpenSim::Geometry& geometry) {
    GeometrySpec spec;
    spec.name = geometry.getName();
    spec.color = toColor(geometry.get_Appearance());
    spec.scale = geometryScale(geometry);

    if(const auto* mesh = dynamic_cast<const OpenSim::Mesh*>(&geometry)) {
        spec.mesh = resolveMeshFile(cache, model, mesh->getGeometryFilename());
    } else if(const auto* sphere = dynamic_cast<const OpenSim::Sphere*>(&geometry)) {
        spec.mesh = sphereMesh(cache);
        spec.scale *= Vector3{Float(sphere->get_radius())};
    } else if(const auto* cylinder = dynamic_cast<const OpenSim::Cylinder*>(&geometry)) {
        spec.mesh = cylinderMesh(cache);
        spec.scale *= Vector3{Float(cylinder->get_radius()), Float(cylinder->get_half_height()), Float(cylinder->get_radius())};
    } else if(const auto* brick = dynamic_cast<const OpenSim::Brick*>(&geometry)) {
        spec.mesh = brickMesh(cache);
        const SimTK::Vec3 half = brick->get_half_lengths();
        spec.scale *= Vector3{Float(half[0]), Float(half[1]), Float(half[2])};
    } else if(const auto* ellipsoid = dynamic_cast<const OpenSim::Ellipsoid*>(&geometry)) {
        spec.mesh = sphereMesh(cache);
        const SimTK::Vec3 radii = ellipsoid->get_radii();
        spec.scale *= Vector3{Float(radii[0]), Float(radii[1]), Float(radii[2])};
    } else {
        throw std::runtime_error("Unsupported geometry type: " + geometry.getConcreteClassName());
    }

    return spec;
}

Matrix4 frameMatrix(const Magnum::Vector3& translation, const Magnum::Quaternion& rotation) {
    return Matrix4::from(rotation.toMatrix(), translation);
}

std::vector<std::string> geometrySearchDirsForModel(const std::string& modelPath,
                                                    const std::vector<std::string>& explicitDirs) {
    namespace fs = std::filesystem;

    std::set<std::string> unique;
    std::vector<std::string> ordered;
    const auto addIfDirectory = [&](const fs::path& path) {
        if(path.empty()) return;
        std::error_code ec;
        if(!fs::is_directory(path, ec)) return;

        const std::string normalized = fs::weakly_canonical(path, ec).string();
        const std::string& key = ec ? path.string() : normalized;
        if(unique.insert(key).second) ordered.push_back(key);
    };

    for(const std::string& dir: explicitDirs) addIfDirectory(fs::path{dir});

    const fs::path modelDir = fs::absolute(fs::path{modelPath}).parent_path();
    addIfDirectory(modelDir);
    addIfDirectory(modelDir/"Geometry");
    if(!modelDir.empty()) addIfDirectory(modelDir.parent_path()/"Geometry");

    #ifdef OSIM_VIEWER_SOURCE_DIR
    addIfDirectory(fs::path{OSIM_VIEWER_SOURCE_DIR}/"Geometry");
    #endif

    return ordered;
}

}

LoadedScene loadScene(const std::string& modelPath,
                      const std::string& motionPath,
                      const std::vector<std::string>& geometryDirs) {
    for(const std::string& dir: geometrySearchDirsForModel(modelPath, geometryDirs)) {
        OpenSim::ModelVisualizer::addDirToGeometrySearchPaths(dir);
    }

    OpenSim::Model model{modelPath};
    auto& coords = model.updCoordinateSet();
    for(int i = 0; i != coords.getSize(); ++i) coords.get(i).setDefaultLocked(false);
    model.initSystem();

    OpenSim::Storage motion{motionPath};
    if(motion.isInDegrees()) model.getSimbodyEngine().convertDegreesToRadians(motion);

    OpenSim::Storage motionAsStates;
    model.formStateStorage(motion, motionAsStates, false);
    const OpenSim::StatesTrajectory trajectory = OpenSim::StatesTrajectory::createFromStatesStorage(model, motionAsStates, true, true);

    MeshCache meshCache;
    LoadedScene scene;
    scene.modelName = model.getName();
    scene.motionName = motion.getName().empty() ? motionPath : motion.getName();

    std::vector<const OpenSim::PhysicalFrame*> framePtrs;
    for(const OpenSim::PhysicalFrame& frame: model.getComponentList<OpenSim::PhysicalFrame>()) {
        if(frame.getProperty_attached_geometry().size() == 0) continue;

        FrameTrack track;
        track.name = frame.getName();
        for(int i = 0; i != frame.getProperty_attached_geometry().size(); ++i) {
            const OpenSim::Geometry& geometry = frame.get_attached_geometry(i);
            if(!geometry.get_Appearance().get_visible()) continue;
            try {
                track.geometries.push_back(makeGeometrySpec(meshCache, model, geometry));
            } catch(const std::exception& error) {
                std::cerr << "Warning: skipping geometry '" << geometry.getName()
                          << "' on frame '" << frame.getName() << "': "
                          << error.what() << '\n';
                continue;
            }
        }
        if(track.geometries.empty()) continue;

        framePtrs.push_back(&frame);
        scene.frames.push_back(std::move(track));
    }

    std::vector<const OpenSim::Marker*> markerPtrs;
    const OpenSim::MarkerSet& markers = model.getMarkerSet();
    for(int i = 0; i != markers.getSize(); ++i) {
        const auto& marker = markers.get(i);
        MarkerTrack track;
        track.name = marker.getName();
        track.color = {0.92f, 0.2f, 0.2f, 1.0f};
        markerPtrs.push_back(&marker);
        scene.markers.push_back(std::move(track));
    }

    scene.times.reserve(trajectory.getSize());
    for(FrameTrack& frame: scene.frames) {
        frame.translations.reserve(trajectory.getSize());
        frame.rotations.reserve(trajectory.getSize());
    }
    for(MarkerTrack& marker: scene.markers) marker.positions.reserve(trajectory.getSize());

    scene.initialBounds.min() = Vector3{std::numeric_limits<Float>::max()};
    scene.initialBounds.max() = Vector3{-std::numeric_limits<Float>::max()};

    for(std::size_t sample = 0; sample != trajectory.getSize(); ++sample) {
        const SimTK::State& state = trajectory[sample];
        scene.times.push_back(Float(state.getTime()));
        model.realizeDynamics(state);

        for(std::size_t i = 0; i != framePtrs.size(); ++i) {
            const SimTK::Transform xform = framePtrs[i]->getTransformInGround(state);
            const Vector3 translation = toVector3(xform.T());
            const Quaternion rotation = toQuaternion(xform.R());
            scene.frames[i].translations.push_back(translation);
            scene.frames[i].rotations.push_back(rotation);

            if(sample == 0) {
                const Matrix4 world = frameMatrix(translation, rotation);
                for(const GeometrySpec& geometry: scene.frames[i].geometries) {
                    const Range3D scaled = scaledBounds(geometry.mesh->bounds, geometry.scale);
                    scene.initialBounds = Math::join(scene.initialBounds, transformBounds(scaled, world));
                }
            }
        }

        for(std::size_t i = 0; i != markerPtrs.size(); ++i) {
            const Vector3 position = toVector3(markerPtrs[i]->getLocationInGround(state));
            scene.markers[i].positions.push_back(position);
            if(sample == 0) scene.initialBounds = Math::join(scene.initialBounds, Range3D::fromSize(position, {}));
        }
    }

    if(scene.frames.empty() && scene.markers.empty()) {
        throw std::runtime_error("The model did not produce any renderable body geometry or markers.");
    }

    return scene;
}

}
