#include "ViewerApp.hpp"

#include <SimTKcommon/internal/PolygonalMesh.h>

#include <Corrade/Containers/ArrayViewStl.h>

#include <Magnum/GL/DefaultFramebuffer.h>
#include <Magnum/GL/Renderer.h>

#include <imgui.h>

#include <algorithm>
#include <limits>
#include <stdexcept>
#include <unordered_map>

using namespace Magnum;
using namespace Math::Literals;

namespace osim_viewer {

namespace {

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

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});
    cpu->indices.reserve(std::size_t(mesh.getNumFaces())*3);
    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 normal = Math::cross(positions[b] - positions[a], positions[c] - positions[a]);
            normals[a] += normal;
            normals[b] += normal;
            normals[c] += normal;
            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;
}

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

}

ViewerApp::MeshDrawable::MeshDrawable(Object3D& object,
                                      Shaders::PhongGL& shader,
                                      GpuMesh& mesh,
                                      Color4 color,
                                      DrawableGroup3D& drawables):
    Drawable3D{object, &drawables},
    _shader{shader},
    _mesh{mesh},
    _color{color} {}

void ViewerApp::MeshDrawable::draw(const Matrix4& transformation, Camera3D& camera) {
    _shader
        .setLightPositions({{4.0f, 6.0f, 4.0f, 0.0f}})
        .setAmbientColor(Color4{_color.rgb()*0.35f, _color.a()})
        .setDiffuseColor(_color)
        .setTransformationMatrix(transformation)
        .setNormalMatrix(transformation.normalMatrix())
        .setProjectionMatrix(camera.projectionMatrix())
        .draw(_mesh.mesh);
}

ViewerApp::ViewerApp(const Arguments& arguments, CliOptions options):
    Platform::Application{arguments, Configuration{}.setTitle("osim-magnum-viewer").setWindowFlags(Configuration::WindowFlag::Resizable)},
    _options{std::move(options)},
    _sceneData{loadScene(_options.modelPath, _options.motionPath, _options.geometryDirs)},
    _shader{Shaders::PhongGL{}},
    _isPlaying{!_options.startPaused},
    _playbackSpeed{_options.initialSpeed} {

    _imgui = ImGuiIntegration::Context{Vector2{windowSize()} / dpiScaling(), windowSize(), framebufferSize()};

    GL::Renderer::enable(GL::Renderer::Feature::DepthTest);
    GL::Renderer::enable(GL::Renderer::Feature::FaceCulling);
    GL::Renderer::setBlendEquation(GL::Renderer::BlendEquation::Add, GL::Renderer::BlendEquation::Add);
    GL::Renderer::setBlendFunction(GL::Renderer::BlendFunction::SourceAlpha, GL::Renderer::BlendFunction::OneMinusSourceAlpha);
    GL::Renderer::setClearColor(0x1e2329_rgbf);

    _camera.setAspectRatioPolicy(SceneGraph::AspectRatioPolicy::Extend);
    _camera.setProjectionMatrix(Matrix4::perspectiveProjection(35.0_degf, Vector2{windowSize()}.aspectRatio(), 0.01f, 500.0f))
        .setViewport(GL::defaultFramebuffer.viewport().size());

    _orbit.setViewport(windowSize(), framebufferSize());
    const Range3D bounds = _sceneData.initialBounds;
    const Vector3 center = (bounds.min() + bounds.max())*0.5f;
    const float radius = Math::max((bounds.max() - bounds.min()).max(), 0.5f);
    _orbit.setSceneBounds(center, radius);
    _cameraObject.setTransformation(_orbit.cameraTransform());
    if(!_sceneData.times.empty()) _currentTime = _sceneData.times.front();

    createSceneObjects();
    updateSceneAtCurrentTime();

    _lastTick = std::chrono::steady_clock::now();
    setMinimalLoopPeriod(16.0_msec);
}

std::shared_ptr<ViewerApp::GpuMesh> ViewerApp::uploadMesh(const CpuMesh& mesh) const {
    auto gpu = std::make_shared<GpuMesh>();
    gpu->bounds = mesh.bounds;
    gpu->vertexBuffer.setData(Containers::arrayView(mesh.vertices));
    gpu->indexBuffer.setData(Containers::arrayView(mesh.indices));
    gpu->mesh.setCount(Int(mesh.indices.size()))
        .setPrimitive(GL::MeshPrimitive::Triangles)
        .addVertexBuffer(gpu->vertexBuffer, 0,
            Shaders::PhongGL::Position{},
            Shaders::PhongGL::Normal{})
        .setIndexBuffer(gpu->indexBuffer, 0, GL::MeshIndexType::UnsignedInt);
    return gpu;
}

void ViewerApp::createSceneObjects() {
    std::unordered_map<const CpuMesh*, std::shared_ptr<GpuMesh>> meshMap;
    auto markerCpu = buildCpuMesh(SimTK::PolygonalMesh::createSphereMesh(1.0, 1));
    _markerMesh = uploadMesh(*markerCpu);
    _gpuMeshes.push_back(_markerMesh);

    for(const FrameTrack& frameTrack: _sceneData.frames) {
        auto frameObject = std::make_unique<Object3D>(&_scene);
        SceneFrame frame;
        frame.object = frameObject.get();

        for(const GeometrySpec& geometry: frameTrack.geometries) {
            std::shared_ptr<GpuMesh> gpuMesh;
            if(const auto found = meshMap.find(geometry.mesh.get()); found != meshMap.end()) {
                gpuMesh = found->second;
            } else {
                gpuMesh = uploadMesh(*geometry.mesh);
                meshMap.emplace(geometry.mesh.get(), gpuMesh);
                _gpuMeshes.push_back(gpuMesh);
            }

            auto geometryObject = std::make_unique<Object3D>(frame.object);
            geometryObject->setTransformation(Matrix4::scaling(geometry.scale));
            auto drawable = std::make_unique<MeshDrawable>(*geometryObject, _shader, *gpuMesh, geometry.color, _drawables);

            frame.geometries.push_back({geometryObject.get(), geometry.color});
            _objectStorage.push_back(std::move(geometryObject));
            _drawableStorage.push_back(std::move(drawable));
        }

        _frames.push_back(frame);
        _objectStorage.push_back(std::move(frameObject));
    }

    constexpr Float MarkerRadius = 0.01f;
    for(const MarkerTrack& markerTrack: _sceneData.markers) {
        auto markerObject = std::make_unique<Object3D>(&_scene);
        markerObject->setTransformation(Matrix4::scaling(Vector3{MarkerRadius}));
        auto drawable = std::make_unique<MeshDrawable>(*markerObject, _shader, *_markerMesh, markerTrack.color, _drawables);
        _markers.push_back({markerObject.get(), markerTrack.color});
        _objectStorage.push_back(std::move(markerObject));
        _drawableStorage.push_back(std::move(drawable));
    }
}

void ViewerApp::updatePlayback() {
    const auto now = std::chrono::steady_clock::now();
    const float delta = std::chrono::duration<float>(now - _lastTick).count();
    _lastTick = now;

    if(!_isPlaying || _sceneData.times.empty()) return;

    _currentTime += delta*_playbackSpeed;
    const float endTime = _sceneData.times.back();
    if(_currentTime >= endTime) {
        _currentTime = endTime;
        _isPlaying = false;
    }
}

std::size_t ViewerApp::sampleIndexForTime(const float time) const {
    if(_sceneData.times.size() < 2) return 0;
    auto it = std::lower_bound(_sceneData.times.begin(), _sceneData.times.end(), time);
    if(it == _sceneData.times.begin()) return 0;
    if(it == _sceneData.times.end()) return _sceneData.times.size() - 2;
    return std::size_t(std::distance(_sceneData.times.begin(), it) - 1);
}

void ViewerApp::updateSceneAtCurrentTime() {
    if(_sceneData.times.empty()) return;

    const std::size_t sample = sampleIndexForTime(_currentTime);
    const std::size_t next = Math::min(sample + 1, _sceneData.times.size() - 1);
    const float t0 = _sceneData.times[sample];
    const float t1 = _sceneData.times[next];
    const float alpha = (next == sample || t1 <= t0) ? 0.0f : (_currentTime - t0)/(t1 - t0);

    for(std::size_t i = 0; i != _frames.size(); ++i) {
        const Vector3 translation = Math::lerp(_sceneData.frames[i].translations[sample], _sceneData.frames[i].translations[next], alpha);
        const Quaternion rotation = Math::slerp(_sceneData.frames[i].rotations[sample], _sceneData.frames[i].rotations[next], alpha).normalized();
        _frames[i].object->setTransformation(makeFrameMatrix(translation, rotation));
    }

    for(std::size_t i = 0; i != _markers.size(); ++i) {
        const Vector3 position = Math::lerp(_sceneData.markers[i].positions[sample], _sceneData.markers[i].positions[next], alpha);
        _markers[i].object->setTransformation(Matrix4::translation(position)*Matrix4::scaling(Vector3{0.01f}));
    }
}

void ViewerApp::drawUi() {
    ImGui::SetNextWindowPos(ImVec2(12.0f, 12.0f), ImGuiCond_Once);
    ImGui::SetNextWindowSize(ImVec2(360.0f, 0.0f), ImGuiCond_Once);
    ImGui::Begin("Playback");
    ImGui::TextUnformatted(_sceneData.modelName.c_str());
    ImGui::TextUnformatted(_sceneData.motionName.c_str());
    if(ImGui::Button(_isPlaying ? "Pause" : "Play")) {
        if(!_isPlaying && !_sceneData.times.empty() && _currentTime >= _sceneData.times.back()) _currentTime = 0.0f;
        _isPlaying = !_isPlaying;
    }
    ImGui::SameLine();
    if(ImGui::Button("Rewind")) {
        _currentTime = 0.0f;
        _isPlaying = false;
    }
    ImGui::SliderFloat("Speed", &_playbackSpeed, 0.1f, 4.0f, "%.2fx");
    ImGui::Text("Time %.3f / %.3f", _currentTime, _sceneData.times.empty() ? 0.0f : _sceneData.times.back());
    ImGui::Text("Frames %zu  Markers %zu", _sceneData.frames.size(), _sceneData.markers.size());
    ImGui::TextUnformatted("LMB orbit  RMB/MMB pan  wheel zoom  F frame");
    ImGui::End();
}

void ViewerApp::drawEvent() {
    updatePlayback();
    updateSceneAtCurrentTime();
    _cameraObject.setTransformation(_orbit.cameraTransform());

    GL::defaultFramebuffer.clear(GL::FramebufferClear::Color | GL::FramebufferClear::Depth);
    _camera.draw(_drawables);

    _imgui.newFrame();
    if(ImGui::GetIO().WantTextInput && !isTextInputActive()) startTextInput();
    else if(!ImGui::GetIO().WantTextInput && isTextInputActive()) stopTextInput();

    drawUi();
    _imgui.updateApplicationCursor(*this);
    GL::Renderer::enable(GL::Renderer::Feature::Blending);
    GL::Renderer::enable(GL::Renderer::Feature::ScissorTest);
    GL::Renderer::disable(GL::Renderer::Feature::FaceCulling);
    GL::Renderer::disable(GL::Renderer::Feature::DepthTest);
    _imgui.drawFrame();
    GL::Renderer::enable(GL::Renderer::Feature::DepthTest);
    GL::Renderer::enable(GL::Renderer::Feature::FaceCulling);
    GL::Renderer::disable(GL::Renderer::Feature::ScissorTest);
    GL::Renderer::disable(GL::Renderer::Feature::Blending);

    swapBuffers();
    redraw();
}

void ViewerApp::viewportEvent(ViewportEvent& event) {
    GL::defaultFramebuffer.setViewport({{}, event.framebufferSize()});
    _imgui.relayout(Vector2{event.windowSize()} / event.dpiScaling(), event.windowSize(), event.framebufferSize());
    _camera.setProjectionMatrix(Matrix4::perspectiveProjection(35.0_degf, Vector2{event.windowSize()}.aspectRatio(), 0.01f, 500.0f))
        .setViewport(event.framebufferSize());
    _orbit.setViewport(event.windowSize(), event.framebufferSize());
}

void ViewerApp::keyPressEvent(KeyEvent& event) {
    if(_imgui.handleKeyPressEvent(event)) return;

    switch(event.key()) {
        case Key::Space:
            if(!_isPlaying && !_sceneData.times.empty() && _currentTime >= _sceneData.times.back()) _currentTime = 0.0f;
            _isPlaying = !_isPlaying;
            event.setAccepted();
            return;
        case Key::R:
            _currentTime = 0.0f;
            _isPlaying = false;
            event.setAccepted();
            return;
        case Key::F: {
            const Vector3 center = (_sceneData.initialBounds.min() + _sceneData.initialBounds.max())*0.5f;
            const float radius = Math::max((_sceneData.initialBounds.max() - _sceneData.initialBounds.min()).max(), 0.5f);
            _orbit.setSceneBounds(center, radius);
            event.setAccepted();
            return;
        }
        case Key::LeftBracket:
            _playbackSpeed = Math::max(_playbackSpeed*0.5f, 0.1f);
            event.setAccepted();
            return;
        case Key::RightBracket:
            _playbackSpeed = Math::min(_playbackSpeed*2.0f, 4.0f);
            event.setAccepted();
            return;
        default:
            break;
    }
}

void ViewerApp::keyReleaseEvent(KeyEvent& event) {
    if(_imgui.handleKeyReleaseEvent(event)) return;
}

void ViewerApp::pointerPressEvent(PointerEvent& event) {
    if(_imgui.handlePointerPressEvent(event)) return;
    _lastPointerPosition = event.position();
    if(event.pointer() == Pointer::MouseLeft) _dragMode = DragMode::Orbit;
    else if(event.pointer() == Pointer::MouseRight || event.pointer() == Pointer::MouseMiddle) _dragMode = DragMode::Pan;
    event.setAccepted();
}

void ViewerApp::pointerReleaseEvent(PointerEvent& event) {
    if(_imgui.handlePointerReleaseEvent(event)) return;
    _dragMode = DragMode::None;
    _lastPointerPosition.reset();
}

void ViewerApp::pointerMoveEvent(PointerMoveEvent& event) {
    if(_imgui.handlePointerMoveEvent(event)) return;
    if(_dragMode == DragMode::None || !_lastPointerPosition) return;

    const Vector2 delta = event.position() - *_lastPointerPosition;
    _lastPointerPosition = event.position();
    if(_dragMode == DragMode::Orbit) _orbit.orbit(delta);
    else if(_dragMode == DragMode::Pan) _orbit.pan(delta);
    event.setAccepted();
}

void ViewerApp::scrollEvent(ScrollEvent& event) {
    if(_imgui.handleScrollEvent(event)) {
        event.setAccepted();
        return;
    }
    _orbit.zoom(event.offset().y());
    event.setAccepted();
}

void ViewerApp::textInputEvent(TextInputEvent& event) {
    if(_imgui.handleTextInputEvent(event)) return;
}

}
