#include "SDL.h"
#include "SDL_image.h"
#include "SDL_ttf.h"
#include <iostream>
#include <assert.h>
#include "shared/path_creator.hpp"
#include "shared/tdmap.hpp"
#include "shared/sizes.h"
#include <map>
#include <cmath>
#include <set>
#include <boost/foreach.hpp>
#include "graphics/graphics_path.hpp"
#include <stdlib.h>

#define TOWER_HOVER_SURFACE "tower_hover"
#define SIDEBAR_OFFSET 1

class GUI
{
    private:
        class Animation
        {
            private:
                size_t cycle_index;

            public:
                GUI * gui;
                SDL_Texture * texture;
                Coordinate position;
                Coordinate from;
                Coordinate to;
                bool randomize;
                std::vector<Coordinate> cycles;
                Coordinate s_size;
                size_t frame_rate;
                int direction;

                Animation(GUI * gui, SDL_Texture * texture, Coordinate position, Coordinate from, Coordinate to, bool randomize, std::vector<Coordinate>& cycles, Coordinate s_size, size_t frame_rate, int direction) :
                    gui(gui), texture(texture), position(position), from(from), to(to), randomize(randomize), cycles(cycles), s_size(s_size), frame_rate(frame_rate), direction(direction)
                    {
                        cycle_index = 0;

                    }

                inline Coordinate get_extra_coord_displacement(int hdiff, int vdiff)
                {
                    int extra_x = hdiff == 0 ? 0 : hdiff/std::abs(hdiff);
                    int extra_y = vdiff == 0 ? 0 : vdiff/std::abs(vdiff);
                    return Coordinate(extra_x, extra_y);
                }

               inline Coordinate get_animation_diff(Coordinate from, Coordinate to)
                {
                    int hdiff = to.x - from.x;
                    int vdiff = to.y - from.y;
                    hdiff = hdiff/ANIMATION_CONSTANT;
                    vdiff = vdiff/ANIMATION_CONSTANT;
                    return Coordinate(hdiff, vdiff);
                }

                void add_to_diff_coords(Coordinate game_coord, Coordinate extras)
                {
                    int extra_x = extras.x;
                    int extra_y = extras.y;
                    Coordinate diagonal_cord = Coordinate(game_coord.x + extra_x, game_coord.y + extra_y);
                    Coordinate horizontal_cord = Coordinate(game_coord.x, game_coord.y + extra_y);
                    Coordinate vertical_cord = Coordinate(game_coord.x + extra_x, game_coord.y);
                    gui->diff_coords.push_back(gui->game_to_screen_coord(diagonal_cord));
                    gui->diff_coords.push_back(gui->game_to_screen_coord(horizontal_cord));
                    gui->diff_coords.push_back(gui->game_to_screen_coord(vertical_cord));
                }

                void animate()
                {
                    if(direction == -1)
                        animate(cycle_index);
                    else
                        animate(direction * frame_rate);
                }

                void animate(size_t specific_cycle_index)
                {
                    cycle_index = specific_cycle_index;
                    Coordinate diffs = get_animation_diff(from, to);
                    Coordinate game_coord = gui->screen_to_game_coord(position);
                    add_to_diff_coords(game_coord, get_extra_coord_displacement(diffs.x, diffs.y));
                    position.x += diffs.x;
                    position.y += diffs.y;
                    gui->render_texture(texture, gui->ren, position, cycles[cycle_index++/frame_rate], s_size);
                    if (cycle_index/frame_rate >= cycles.size())
                    {
                        cycle_index = 0;
                    }
                }
        };

        //tiles
        int numrows;
        int numcols;
        int row_width;
        int row_height;

        std::vector<Path> paths;
        TDMap * map;
        SDL_Texture * tile;
        SDL_Texture * background;
        SDL_Window *win;
        SDL_Renderer * ren;
        std::vector<Animation> animations;
        std::vector<Coordinate> diff_coords;
        std::map<std::string, SDL_Texture *> loaded_textures;
        typedef std::pair<std::string, bool> file_string_colored;
        std::map<std::string, SDL_Surface *> loaded_surfaces;
        std::map<file_string_colored, SDL_Texture*> loaded_sprite_textures;
        std::map<Coordinate, bool> colors;
        SDL_Thread * event_thread;
        bool quit;
        SDL_Texture * tile_hover;
        SDL_Texture * sidebar;
        SDL_mutex *lock;
        SDL_cond *cond;
        TTF_Font * font;
        int previous_score;
        int actual_width;

        void logSDLError(std::ostream &os, const std::string &msg)
        {
            os << msg << " error: " << SDL_GetError() << std::endl;
        }

        SDL_Surface * load_surface(const std::string & file)
        {
            if(loaded_surfaces.find(file) != loaded_surfaces.end())
            {
                return loaded_surfaces[file];
            }
            else
            {
                SDL_Surface * surface = IMG_Load((GRAPHICS_PATH + file).c_str());
                loaded_surfaces[file] = surface;
                return surface;
            }

        }

        SDL_Texture * load_sprite_texture(const std::string &file, bool colored)
        {
            file_string_colored f = std::make_pair(file, colored);
            if(loaded_sprite_textures.find(f) != loaded_sprite_textures.end())
            {
                return loaded_sprite_textures[f];
            }
            else
            {
                SDL_Surface * surface = load_surface(file);
                SDL_Texture * texture = SDL_CreateTextureFromSurface(ren, surface);
                if(colored)
                    SDL_SetTextureColorMod(texture, 255, 0, 0);
                loaded_sprite_textures[f] = texture;
                return texture;
            }
        }

        SDL_Texture* load_texture(const std::string &file, SDL_Renderer *ren)
        {
            SDL_Texture *texture;

            if(loaded_textures.find(file) != loaded_textures.end())
            {
                texture = loaded_textures[file];
            }
            else
            {
                SDL_Surface * surface = load_surface(file);
                if (surface == nullptr)
                {
                    logSDLError(std::cout, "LoadTexture");
                }
                texture = SDL_CreateTextureFromSurface(ren, surface);
                if (texture == nullptr)
                {
                    logSDLError(std::cout, "LoadTexture");
                }
                loaded_textures[file] = texture;
            }
            return texture;
        }

        void render_texture(SDL_Texture * tex, SDL_Renderer * ren, Coordinate screen_location, Coordinate sprite_coordinate, Coordinate sizes)
        {
            SDL_Rect renderQuad = {sprite_coordinate.x * sizes.x, sprite_coordinate.y * sizes.y , sizes.x, sizes.y};
            render_texture(tex, ren, &renderQuad, screen_location);
        }

        void render_texture(SDL_Texture * tex, SDL_Renderer * ren, Coordinate coordinate)
        {
            render_texture(tex, ren, nullptr, coordinate);
        }

        void render_texture(SDL_Texture * tex, SDL_Renderer * ren, Coordinate coordinate, int width, int height)
        {
            render_texture(tex, ren, nullptr, coordinate.x, coordinate.y, width, height);
        }

        void render_texture(SDL_Texture * tex, SDL_Renderer * ren, SDL_Rect * rect, Coordinate coordinate)
        {
            render_texture(tex, ren, rect, coordinate.x, coordinate.y, row_width, row_height);
        }

        void render_texture(SDL_Texture *tex, SDL_Renderer *ren, SDL_Rect* clip_rect, int x, int y, int rw, int rh)
        {
            SDL_Rect dst;
            dst.x = x;
            dst.y = y;
            dst.w = rw;
            dst.h = rh;
            SDL_RenderCopy(ren, tex, clip_rect, &dst);
        }

        Coordinate screen_to_game_coord(Coordinate c)
        {
            return Coordinate(c.x / row_width, c.y / row_height);
        }

        Coordinate game_to_screen_coord(Coordinate game_coord)
        {
            return Coordinate(game_coord.x * row_width, game_coord.y * row_height);
        }

        Coordinate game_to_screen_coord(int x, int y)
        {
            return game_to_screen_coord(Coordinate(x, y));
        }

        void render_game_screen(Coordinate game_coord, bool colorize)
        {

            Coordinate screen_cord = game_to_screen_coord(game_coord);

            if(false)
            {
                SDL_SetTextureColorMod(background, 255, 128, 128);
                SDL_SetTextureColorMod(tile, 255, 128, 128);
            }
            else
            {
                SDL_SetTextureColorMod(background, 255, 255, 255);
                SDL_SetTextureColorMod(tile, 255, 255, 255);
            }

            render_texture(background, ren, screen_cord);

            BOOST_FOREACH(Path & path, paths)
            {
                if(path.in(game_coord))
                {
                    render_texture(tile, ren, screen_cord);
                }
            }

            Tower * tower = map->get_tower_at(game_coord);

            SDL_SetRenderTarget(ren, background);

            if(tower != nullptr)
            {
                SDL_Texture * texture = load_texture(tower->get_image_string(), ren);
                if(texture != nullptr)
                {
                    render_texture(texture, ren, screen_cord);
                }
            }
            SDL_SetRenderTarget(ren, nullptr);
        }

        void re_render()
        {
            fill_screen_tiles(false);
        }

        void fill_screen_tiles(bool colorize)
        {
            SDL_LockMutex(lock);

            for(int i = 0; i < numrows; i++)
            {
                for(int j = 0; j < numcols; j++)
                {
                    render_game_screen(Coordinate(i, j), colorize);
                }
            }

            SDL_UnlockMutex(lock);

        }

        void render_sidebar()
        {
            for(int i = numrows; i < numrows + SIDEBAR_OFFSET; i++)
            {
                for(int j = 0; j < numcols; j++)
                {
                    render_sidebar(game_to_screen_coord(i, j));
                }
            }
        }

        void render_sidebar(Coordinate screen_cord)
        {
            render_texture(sidebar, ren, screen_cord, actual_width, row_height);
        }

        Tower * last_tower;

        void handle_mouse_over_at(Coordinate screen_cord)
        {

            Tower * tower = map->get_tower_at(screen_to_game_coord(screen_cord));
            SDL_LockMutex(lock);
            if(tower != nullptr)
            {
                for(int i = 0; i < numrows; i++)
                {
                    for(int j = 0; j < numcols; j++)
                    {
                        if(tower->can_attack(Coordinate(i, j)))
                        {

                            colors[Coordinate(i, j)] = true;
                        }
                    }
                }
            }

            else
            {
                colors.clear();
            }

            SDL_UnlockMutex(lock);

            last_tower = tower;

        }

        static int handle_event(void * ptr)
        {
             return ((GUI*)ptr)->handle_event();
        }

        int handle_event()
        {

            SDL_Event event;

            while(!quit)
            {

                /* Poll for events */
                while(SDL_PollEvent( &event ))
                {

                    switch( event.type )
                    {
                        /* SDL_QUIT event (window close) */
                        case SDL_QUIT:
                            std::cout << "Tried to quit!" << std::endl;
                            quit = true;
                            break;

                        case SDL_MOUSEMOTION:
                            handle_mouse_over_at(Coordinate(event.motion.x, event.motion.y));
                            break;

                        default:
                            break;
                    }

                }

            }

            SDL_Quit();
            return 0;
        }

        bool init()
        {
            return SDL_Init(SDL_INIT_EVERYTHING) == 0;
        }

        bool create_window()
        {
            win = SDL_CreateWindow("Tower Defense",
                    SDL_WINDOWPOS_UNDEFINED, SDL_WINDOWPOS_UNDEFINED, DEFAULT_WIDTH, DEFAULT_HEIGHT, SDL_WINDOW_SHOWN);

            return win != nullptr;
        }

        bool create_renderer()
        {
            ren = SDL_CreateRenderer(win, -1, SDL_RENDERER_PRESENTVSYNC);
            return ren != nullptr;
        }

        bool create_textures()
        {

            background = load_texture("grass_night.jpg", ren);
            tile = load_texture("tile2.png", ren);
            sidebar = load_texture("tile.png", ren);
            SDL_Surface * tile_hover_s = SDL_CreateRGBSurface(0, row_width, row_height, 32, 0, 0, 0, 0);
            SDL_FillRect(tile_hover_s, nullptr, SDL_MapRGB(tile_hover_s->format, 255, 0, 0));
            tile_hover = SDL_CreateTextureFromSurface(ren, tile_hover_s);
            loaded_textures[TOWER_HOVER_SURFACE] = tile_hover;
            return background != nullptr || tile != nullptr || tile_hover != nullptr;
        }

        bool create_font()
        {
            if( TTF_Init() == -1 )
            {
                return false;
            }

            std::string res_path = GRAPHICS_PATH;
            std::string extra = "Lato-Regular.ttf";

            font = TTF_OpenFont((res_path + extra).c_str() , 28 );

            if(font == nullptr)
            {
                return false;
            }

            previous_score = -1;

            return true;
        }

        SDL_Texture * get_font_texture(const std::string &message, SDL_Color color)
        {

            if(loaded_textures.find(message) != loaded_textures.end())
                return loaded_textures[message];

            SDL_Surface * surf = TTF_RenderText_Solid(font, message.c_str(), color);
            if (surf == nullptr)
            {
                logSDLError(std::cout, "TTF_RenderText");
                return nullptr;
            }
            SDL_Texture *texture = SDL_CreateTextureFromSurface(ren, surf);

            if (texture == nullptr)
            {
                logSDLError(std::cout, "CreateTexture");
            }

            loaded_textures[message] = texture;

            return texture;
        }

    public:

       bool has_quit_game()
        {
            return quit;
        }

        GUI(int rows, int columns, std::vector<Path> &paths, TDMap * map)
        {

            assert(rows > 0 && columns > 0 && paths.size() != 0 && map != nullptr);

            actual_width = DEFAULT_WIDTH/rows;
            //row_width = (rows * actual_width - SIDEBAR_OFFSET * actual_width)/rows;
            row_width = actual_width;
            row_height = DEFAULT_HEIGHT/columns;
            numrows = rows;
            numcols = columns;
            this->map = map;
            this->paths = paths;
            lock = SDL_CreateMutex();
            /* Doesnt work with short circuiting in an if... wonder why */

            bool completed;
            completed = init();
            completed = completed && create_window();
            completed = completed && create_renderer();
            completed = completed && create_textures();
            completed = completed && create_font();

            if (!completed)
            {
                logSDLError(std::cout, "Initialization of components");
                SDL_Quit();
                exit(1);
            }

            SDL_Thread * event_thread = SDL_CreateThread(handle_event, "hello", this);

            fill_screen_tiles(false);
            SDL_RenderPresent(ren);
        }

        std::vector<Animation> tower_anims;
        std::vector<Animation> sprite_anims;

        // takes pixel coords.
        // pls take care while using.

        void set_up_tower_animation(SDL_Texture * texture, Coordinate from, Coordinate to, bool randomize, std::vector<Coordinate>& s_locations, Coordinate s_size, size_t frame_rate, int direction)
        {
            set_up_animation(tower_anims, texture, from, to, randomize, s_locations, s_size, frame_rate, direction);
        }

        void set_up_sprite_animation(Sprite * sprite, SDL_Texture * texture, Coordinate from, Coordinate to, bool randomize, std::vector<Coordinate>& s_locations, Coordinate s_size, size_t frame_rate, int direction)
        {
            if(from.x < 0 && randomize)
            {
                Coordinate offsets(((rand() % (row_width + 1)) - (row_width)/2)/2, (rand() % (row_height + 1)) - (row_height)/2);
                sprite->set_random_position(offsets);
            }

            Coordinate offsets = sprite->get_offset();
            from.x += offsets.x;
            from.y += offsets.y;
            to.x += offsets.x;
            to.y += offsets.y;

            set_up_animation(sprite_anims, texture, from, to, randomize, s_locations, s_size, frame_rate, direction);
        }

        void set_up_animation(std::vector<Animation>& which, SDL_Texture * texture, Coordinate from, Coordinate to, bool randomize, std::vector<Coordinate>& s_locations, Coordinate s_size, size_t frame_rate, int direction)
        {
            which.push_back(Animation(this, texture, from, from, to, randomize, s_locations, s_size, frame_rate, direction));
            diff_coords.push_back(from);
            diff_coords.push_back(to);
        }

        void clear_animations()
        {
            re_render();
            tower_anims.clear();
            sprite_anims.clear();
            SDL_LockMutex(lock);
            colors.clear();
            SDL_UnlockMutex(lock);

        }

        void show_sprite_animations()
        {
            re_render();
            BOOST_FOREACH(Animation & animation, sprite_anims)
            {
                animation.animate();
            }
        }

        void show_tower_animations()
        {
            BOOST_FOREACH(Animation & animation, tower_anims)
            {
                animation.animate();
            }
        }

        void Update()
        {

            SDL_Color red = {255, 0, 0, 0};
            int score = map->get_number_of_sprites();
            if(score != previous_score)
            {
                render_sidebar();
                SDL_Texture * tex = get_font_texture(std::to_string(score), red);
                Coordinate right_top = game_to_screen_coord(Coordinate(numrows, 0));
                render_texture(tex, ren, right_top);
            }
            previous_score = score;

            SDL_LockMutex(lock);
            for(int i = 0; i < NUM_ROWS; i++)
            {
                for(int j = 0; j < NUM_COLS; j++)
                {
                    BOOST_FOREACH(Sprite * s, map->get_sprites_at(i,j))
                    {
                        SDL_Texture * sprite_texture = load_sprite_texture(s->get_spritesheet()->get_file_name(), s->is_attacked());
                        Coordinate previous_cord = game_to_screen_coord(s->get_previous_position());
                        Coordinate next_cord = game_to_screen_coord(s->get_coordinate());
                        set_up_sprite_animation(s, sprite_texture, previous_cord, next_cord, true, s->get_sscords(), s->get_spritesheet()->get_box_size(), s->get_spritesheet()->get_frame_rate(), -1);
                    }

                    Tower * tower = map->get_tower_at(i, j);
                    if(tower != nullptr)
                    {
                        Coordinate current_cord = game_to_screen_coord(Coordinate(i, j));
                        SDL_Texture * texture = load_texture(tower->get_spritesheet()->get_file_name(), ren);
                        BOOST_FOREACH(Coordinate c, tower->get_attack_tiles())
                        {
                            colors[c] = true;
                            Coordinate screen_cord = game_to_screen_coord(c);
                            int xdirection = current_cord.x < screen_cord.x ? 2 : current_cord.x == screen_cord.x ? 1 : 0;
                            int ydirection = current_cord.y < screen_cord.y ? 0 : 1;
                            int direction = xdirection + (tower->get_sscords().size()/tower->get_spritesheet()->get_frame_rate() * ydirection);
                            set_up_tower_animation(texture, current_cord, screen_cord, false, tower->get_sscords(), tower->get_spritesheet()->get_box_size(), tower->get_spritesheet()->get_frame_rate(), direction);
                        }
                    }
                }
            }

            SDL_UnlockMutex(lock);
            for(int i = 0; i < ANIMATION_CONSTANT; i++)
            {
                show_sprite_animations();
                show_tower_animations();
                SDL_RenderPresent(ren);
            }
            SDL_RenderPresent(ren);
            clear_animations();
        }
};


#include <iostream>
#include <boost/thread.hpp>
#include <drc_utils/Clock.hpp>
#include <maps/PointDataBuffer.hpp>
#include <maps/SensorDataReceiver.hpp>
#include <bot_frames/bot_frames.h>
#include <bot_param/param_client.h>
#include <ConciseArgs>
#include <bot_lcmgl_client/lcmgl.h>

#include <pointcloud_tools/pointcloud_math.hpp>
#include <pointcloud_tools/pointcloud_vis.hpp>

using namespace std;

struct State {
  boost::shared_ptr<lcm::LCM> mLcm;
  boost::shared_ptr<PointDataBuffer> mPointDataBuffer;
  boost::shared_ptr<SensorDataReceiver> mSensorDataReceiver;
  BotParam* mBotParam;
  BotFrames* mBotFrames;
  bot_lcmgl_t* mLcmGl;

  State() {
    mPointDataBuffer.reset(new PointDataBuffer());
    mSensorDataReceiver.reset(new SensorDataReceiver());
    mLcm.reset(new lcm::LCM());
    drc::Clock::instance()->setLcm(mLcm);
    mSensorDataReceiver->setLcm(mLcm);
    mBotParam = bot_param_new_from_server(mLcm->getUnderlyingLCM(), 0);
    mSensorDataReceiver->setBotParam(mBotParam);
    mBotFrames = bot_frames_get_global(mLcm->getUnderlyingLCM(), mBotParam);
    mLcmGl = bot_lcmgl_init(mLcm->getUnderlyingLCM(), "recent-points");
  }

  ~State() {
    bot_param_destroy(mBotParam);
    bot_lcmgl_destroy(mLcmGl);
  }
};

class DataConsumer {
public:
  DataConsumer(State* iState) {
    mState = iState;
    mPrevAngle_ = -1;
    mPrevTime = -1;
    counter_ =0;

    pc_vis_ = new pointcloud_vis(mState->mLcm->getUnderlyingLCM());

  }

  void operator()() {
    while(true) {
      SensorDataReceiver::PointCloudWithPose data;
      if (mState->mSensorDataReceiver->waitForData(data)) {
        PointDataBuffer::PointSet pointSet;
        pointSet.mTimestamp = data.mTimestamp;
        pointSet.mPoints = data.mPointCloud;
        pointSet.mToLocal = data.mPose;
        mState->mPointDataBuffer->add(pointSet);

        BotTrans b2l;
        bot_frames_get_trans_with_utime(mState->mBotFrames, "body", "local",
                                        data.mTimestamp, &b2l);

        BotTrans transform;
        bot_frames_get_trans_with_utime(mState->mBotFrames, "ROTATING_SCAN", "head",
                                        data.mTimestamp, &transform);
        double currentAngle = 2*acos(transform.rot_quat[0]);
        if (currentAngle > M_PI) {
          currentAngle -= 2*M_PI;
        }
        if (mPrevAngle_*currentAngle < 0) {
          std::cout << "Previous=" << (mPrevAngle_*180/M_PI) <<
            ", Current=" << (currentAngle*180/M_PI) << std::endl;

          maptypes::PointCloud::Ptr cloud =
            mState->mPointDataBuffer->getAsCloud(mPrevTime, data.mTimestamp);

          bot_lcmgl_t* lcmgl = mState->mLcmGl;

          bot_lcmgl_color3f(lcmgl, pc_vis_->colors[counter_*3], pc_vis_->colors[counter_*3+1], pc_vis_->colors[counter_*3+2]);
          for (int i = 0; i < cloud->points.size(); ++i) {
            maptypes::PointCloud::PointType point = cloud->points[i];
            bot_lcmgl_begin(lcmgl, LCMGL_POINTS);
            bot_lcmgl_vertex3f(lcmgl, point.x, point.y, point.z);
            bot_lcmgl_end(lcmgl);
          }
          bot_lcmgl_switch_buffer(lcmgl);

          std::cout << "Published " << cloud->points.size() << " points " << std::endl;
          mPrevTime = data.mTimestamp;

          ostringstream fname;
          fname << "cloud " << b2l.trans_vec[0] << " "
              << b2l.trans_vec[1] << " "
              << b2l.trans_vec[2] << " "
              << b2l.rot_quat[0] << " "
              << b2l.rot_quat[1] << " "
              << b2l.rot_quat[2] << " "
              << b2l.rot_quat[3] << ".pcd";
          pcl::io::savePCDFileASCII (fname.str(), *cloud);
          std::cerr << data.mTimestamp << " =========Clock\n";
          std::cerr << "Saved " << cloud->points.size () << " data points to " << fname.str() << std::endl;
          counter_++;
          if(counter_*3 >= pc_vis_->colors.size() ){
           counter_=0;
          }

          //pc_vis_->pose_collection_to_lcm_from_list(6001, world_to_jointsT); // all joints in world frame

        }

        mPrevAngle_ = currentAngle;

      }
    }
  }

protected:
  State* mState;
  double mPrevAngle_;
  int64_t mPrevTime;
  int counter_;
  pointcloud_vis* pc_vis_;
};

int main(const int iArgc, const char** iArgv) {
  State state;

  // parse arguments
  double mapResolution = 0.02;
  string laserChannel = "ROTATING_SCAN";
  double xDim(10), yDim(10), zDim(10);
  ConciseArgs opt(iArgc, (char**)iArgv);
  opt.add(laserChannel, "l", "laser_channel", "channel for incoming scan data");
  opt.parse();
  state.mSensorDataReceiver->
    addChannel(laserChannel,
               SensorDataReceiver::SensorTypePlanarLidar,
               laserChannel, "local");

  // set up remaining parameters
  state.mSensorDataReceiver->setMaxBufferSize(100);

  // start running data receiver
  state.mSensorDataReceiver->start();

  // start consuming data
  DataConsumer consumer(&state);
  boost::thread thread(consumer);

  // main lcm loop
  while (0 == state.mLcm->handle());

  state.mSensorDataReceiver->stop();
  thread.join();

  return 0;
}


/**
 * MIT License
 *
 * Copyright (c) 2020 IWT Wirtschaft und Technik GmbH
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 * @file box_detector_node.cpp
 * @author Vishnuprasad Prachandabhanu (vishnu.pbhat93@gmail.com)
 */
#include <iostream>
#include <thread>
#include <mutex>

#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2_ros/static_transform_broadcaster.h>
#include <tf2_ros/transform_listener.h>
#include <geometry_msgs/TransformStamped.h>

#include <pcl_ros/point_cloud.h>
#include <pcl/io/impl/synchronized_queue.hpp>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/io/pcd_io.h>
#include <boost/foreach.hpp>
#include <boost/filesystem.hpp>

#include <box_detector/box_detector.hpp>

typedef pcl::PointXYZ PointType;
typedef pcl::PointCloud<PointType> PointCloud;

class Detector{
    private:
        ros::NodeHandle _nh;
        ros::Subscriber _sub;
        ros::Publisher _pub;
        ros::Publisher _pubPose;
        PointCloud::Ptr _cloud;
        tf2_ros::Buffer buffer;
        pcl::SynchronizedQueue<pcl::PointCloud<pcl::PointXYZ>> _queue;

        bool _newCloud = false;
        std::mutex cloud_lock;
        double distance_max, distance_min, per_width, per_height, width, length, height;

        nimbus::BoxDetector<pcl::PointXYZ> * boxDectect;

        Eigen::Matrix<float, 4, 1> centroid, param_norm;
        float yaw = 0;

        tf2_ros::StaticTransformBroadcaster broadCaster;
        geometry_msgs::TransformStamped pose;

        unsigned int yawCounter;

    public:
        Detector(ros::NodeHandle nh): _nh(nh)
        {
            _sub = _nh.subscribe<sensor_msgs::PointCloud2>("/nimbus/pointcloud", 10, boost::bind(&Detector::callback, this, _1));
            _pub = _nh.advertise<PointCloud>("filtered_cloud", 5);
            _pubPose = _nh.advertise<geometry_msgs::TransformStamped>("detected_pose", 10);
            tf2_ros::TransformListener listener(buffer);

            this->boxDectect = new nimbus::BoxDetector<pcl::PointXYZ>(_nh);

            pose.header.frame_id = "camera";
            pose.child_frame_id = "box";
            yawCounter = 0;
        }

        ~Detector(){}

        void callback(const sensor_msgs::PointCloud2::ConstPtr &msg)
        {
            std::lock_guard<std::mutex> lock(cloud_lock);
            _cloud.reset(new PointCloud());
            pcl::PointCloud<pcl::PointXYZI>::Ptr blob (new pcl::PointCloud<pcl::PointXYZI>());
            pcl::PCLPointCloud2 pcl_pc2;
            pcl_conversions::toPCL(*msg, pcl_pc2);
            pcl::fromPCLPointCloud2(pcl_pc2, *blob);
            _cloud->is_dense = false;
            pcl::copyPointCloud(*blob, *_cloud);
            _newCloud = true;
        }

        void updateParm(ros::NodeHandle nh)
        {
            nh.getParam("distance_max", distance_max);
            nh.getParam("distance_min", distance_min);
            nh.getParam("per_width", per_width);
            nh.getParam("per_height", per_height);
            nh.getParam("box_width", width);
            nh.getParam("box_length", length);
            nh.getParam("box_height", height);
        }

        bool groudTruth(const boost::shared_ptr< const pcl::PointCloud<pcl::PointXYZ>> blob, pcl::PointCloud<pcl::PointXYZ> &res)
        {
            std::string model_name = "/grount_truth";
            std::string extention = ".pcd";
            std::string saved_groudtruth;
            int counter = 0;
            std::stringstream dir;
            std::string _working_dir = getenv("HOME");
            dir << _working_dir << "/ros_ws/ground";
            boost::filesystem::path test_dir = dir.str();
            if(!boost::filesystem::exists(test_dir) && !boost::filesystem::is_directory(test_dir))
            {
                ROS_INFO("------------------ Creating folder ---------------------");
                boost::filesystem::create_directory(test_dir);
                ros::Duration(1.0).sleep();
            }
            std::stringstream file;
            file << test_dir.c_str() << model_name << extention;
            if(!boost::filesystem::exists(file.str()))
            {
                ROS_INFO("------------------ For the first time scan the empty table ---------------------");
                ROS_INFO("------------------ If the table is not empty then please empty the table and RESTART ---------------------");
                ros::Duration(1.0).sleep();
                ROS_INFO("                   Saving ground truth.....");
                pcl::io::savePCDFile(file.str(), *blob);
                ROS_INFO("                   Place the box");
                return false;
            }
            pcl::PointCloud<pcl::PointXYZ>::Ptr cloud (new pcl::PointCloud<pcl::PointXYZ>());
            typename pcl::PointCloud<pcl::PointXYZ>::Ptr ground (new pcl::PointCloud<pcl::PointXYZ>());
            pcl::io::loadPCDFile(file.str(), *ground);
            bool model = boxDectect->getBaseModel(ground, blob, height - 0.04, file.str(), *cloud);
            if(!model) return false;
            pcl::copyPointCloud(*cloud, res);
            return true;
        }

        void run()
        {
            ros::spinOnce();
            ros::Rate rate(5);
            while(ros::ok())
            {
                updateParm(this->_nh);
                if(_newCloud)
                {
                    _newCloud = false;
                    PointCloud::Ptr blob (new PointCloud());
                    PointCloud::Ptr rCloud (new PointCloud());
                    PointCloud::Ptr meanCloud (new PointCloud());
                    PointCloud::Ptr cloud (new PointCloud());

                    std::unique_lock<std::mutex> lock(cloud_lock);
                    pcl::copyPointCloud(*_cloud, *blob);
                    lock.unlock();

                    boxDectect->outlineRemover(blob, blob->width, blob->height, per_width, per_height, *rCloud);

                    _queue.enqueue(*rCloud);
                    // Queue sheild
                    if(_queue.size() < 2) continue;

                    boxDectect->meanFilter(_queue, *meanCloud);

                    bool model = groudTruth(meanCloud, *cloud);
                    if(!model) continue;
                    //// Core Operation ////
                    boxDectect->box3DCentroid(cloud, centroid);
                    if(std::isnan(centroid[0])){
                        ROS_ERROR ("Can not find the centroid");
                        continue;
                    }
                    bool calYaw = boxDectect->boxYaw(cloud, width, length, centroid, yaw);
                    ////////////////////////
                    if(calYaw)
                    {
                        // Dump first few values of yaw make it stable
                        // ToDo replace this loop burner
                        // if(yawCounter <= 5)
                        // {
                        //     yawCounter += 1;
                        //     continue;
                        // }
                        // yawCounter = 0;

                        if((yaw * 180)/M_PI > 90) yaw = yaw - M_PI;
                        if((yaw * 180)/M_PI < -90) yaw = yaw + M_PI;
                        ROS_INFO("Yaw :%f", (yaw * 180)/M_PI );
                        pose.header.stamp = ros::Time::now();
                        pose.transform.translation.x = centroid[0];
                        pose.transform.translation.y = centroid[1];
                        pose.transform.translation.z = centroid[2];
                        tf2::Quaternion q;
                        // ToDo Orientation correction instead of -ve in x and y
                        q.setRPY(0,0,yaw);
                        pose.transform.rotation = tf2::toMsg(q);
                        _pubPose.publish(pose);
                        broadCaster.sendTransform(pose);
                    }

                    cloud->header.frame_id = "camera";
                    pcl_conversions::toPCL(ros::Time::now(), cloud->header.stamp);
                    _pub.publish(cloud);
                    ros::spinOnce();
                }else{
                    ros::spinOnce();
                }
                rate.sleep();
            }
        }
};

int main(int argc, char** argv){
    ros::init(argc, argv, "box_detector_node");
    ros::NodeHandle nh("~");
    Detector detector(nh);
    try{
        detector.run();
    }catch(ros::Exception e){
        ROS_ERROR("Runtime error %s", e.what());
    }
    return 0;
}


//==================================================================================================
/**
  Copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
**/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_X86_AVX_SIMD_FUNCTION_IF_ELSE_HPP_INCLUDED
#define BOOST_SIMD_ARCH_X86_AVX_SIMD_FUNCTION_IF_ELSE_HPP_INCLUDED

#include <boost/simd/detail/overload.hpp>
#include <boost/simd/detail/dispatch/meta/as_floating.hpp>
#include <boost/simd/meta/cardinal_of.hpp>
#include <boost/simd/function/bitwise_cast.hpp>
#include <boost/simd/function/genmask.hpp>

namespace boost { namespace simd { namespace ext
{
   namespace bd = boost::dispatch;
   namespace bs = boost::simd;

   BOOST_DISPATCH_OVERLOAD( if_else_
                          , (typename A0, typename A1)
                          , bs::avx_
                          , bs::pack_<logical_<A0>, bs::avx_>
                          , bs::pack_<bd::single_<A1>, bs::avx_>
                          , bs::pack_<bd::single_<A1>, bs::avx_>
                          )
   {
     A1 operator()(A0 const& a0,A1 const& a1,A1 const& a2) const
     {
       return _mm256_blendv_ps(a2, a1, bitwise_cast<A1>(genmask(a0)));
     }
   };

   BOOST_DISPATCH_OVERLOAD( if_else_
                          , (typename A0, typename A1)
                          , bs::avx_
                          , bs::pack_<logical_<A0>, bs::avx_>
                          , bs::pack_<bd::double_<A1>, bs::avx_>
                          , bs::pack_<bd::double_<A1>, bs::avx_>
                          )
   {
     A1 operator()(A0 const& a0,A1 const& a1,A1 const& a2) const
     {
       return _mm256_blendv_pd(a2, a1, bitwise_cast<A1>(genmask(a0)));
     }
   };

   BOOST_DISPATCH_OVERLOAD_IF ( if_else_
                              , (typename A0, typename A1)
                              , (brigand::bool_<bs::cardinal_of<A0>::value <= 8>)
                              , bs::avx_
                              , bs::pack_<logical_<A0>, bs::avx_>
                              , bs::pack_<bd::integer_<A1>, bs::avx_>
                              , bs::pack_<bd::integer_<A1>, bs::avx_>
                              )
   {
     A1 operator()(A0 const& a0,A1 const& a1,A1 const& a2) const
     {
        using f_t= bd::as_floating_t<A1>;
        return bitwise_cast<A1>(if_else(a0, bitwise_cast<f_t>(a1), bitwise_cast<f_t>(a2)));
     }
   };
} } }

#endif


//---------------------------------------------------------------------------//
//  MIT License
//
//  Copyright (c) 2020-2021 Mikhail Komarov <nemo@nil.foundation>
//  Copyright (c) 2020-2021 Nikita Kaskov <nemo@nil.foundation>
//  Copyright (c) 2021-2022 Aleksei Moskvin <alalmoskvin@gmail.com>
//  Copyright (c) 2021 Ilias Khairullin <ilias@nil.foundation>
//
//  Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
//
//  The above copyright notice and this permission notice shall be included in all
//  copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//  SOFTWARE.
//---------------------------------------------------------------------------//

#ifndef CRYPTO3_MERKLE_PROOF_HPP
#define CRYPTO3_MERKLE_PROOF_HPP

#include <algorithm>
#include <vector>

#include <boost/variant.hpp>

#include <nil/crypto3/container/merkle/tree.hpp>

namespace nil {
    namespace crypto3 {
        namespace zk {
            namespace components {
                template<typename, typename, std::size_t>
                struct merkle_proof;
            }    // namespace components
        }        // namespace zk
        namespace marshalling {
            namespace types {
                template<typename, typename>
                struct merkle_proof_marshalling;
            }
        }    // namespace marshalling
        namespace containers {
            namespace detail {
                template<typename NodeType, std::size_t Arity = 2>
                struct merkle_proof_impl {
                    typedef NodeType node_type;
                    typedef typename node_type::hash_type hash_type;

                    constexpr static const std::size_t arity = Arity;

                    constexpr static const std::size_t value_bits = node_type::value_bits;
                    typedef typename node_type::value_type value_type;

                    merkle_proof_impl() : _li(0) {};

                    merkle_proof_impl(const merkle_tree<hash_type, arity> &tree, const std::size_t leaf_idx) {
                        _root = tree.root();
                        _path.resize(tree.row_count() - 1);
                        _li = leaf_idx;

                        typename std::vector<layer_type>::iterator v_itr = _path.begin();
                        std::size_t cur_leaf = leaf_idx;
                        std::size_t row_len = tree.leaves();
                        std::size_t row_begin_idx = 0;
                        while (cur_leaf != tree.size() - 1) {    // while it's not _root
                            std::size_t cur_leaf_pos = cur_leaf % arity;
                            std::size_t cur_leaf_arity_pos = (cur_leaf - row_begin_idx) / arity;
                            std::size_t begin_this_arity = cur_leaf - cur_leaf_pos;
                            typename layer_type::iterator a_itr = v_itr->begin();
                            for (size_t i = 0; i < cur_leaf_pos; ++i, ++begin_this_arity, ++a_itr) {
                                *a_itr = path_element_type(tree[begin_this_arity], i);
                            }
                            for (size_t i = cur_leaf_pos + 1; i < arity; ++i, ++begin_this_arity, ++a_itr) {
                                *a_itr = path_element_type(tree[begin_this_arity + 1], i);
                            }
                            v_itr++;
                            cur_leaf = row_len + row_begin_idx + cur_leaf_arity_pos;
                            row_begin_idx += row_len;
                            row_len /= arity;
                        }
                    }

                    template<typename Hashable>
                    bool validate(const Hashable &a) const {
                        value_type d = crypto3::hash<hash_type>(a);
                        for (auto &it : _path) {
                            accumulator_set<hash_type> acc;
                            size_t i = 0;
                            for (; (i < arity - 1) && i == it[i]._position; ++i) {
                                crypto3::hash<hash_type>(it[i]._hash.begin(), it[i]._hash.end(), acc);
                            }
                            crypto3::hash<hash_type>(d.begin(), d.end(), acc);
                            for (; i < arity - 1; ++i) {
                                crypto3::hash<hash_type>(it[i]._hash.begin(), it[i]._hash.end(), acc);
                            }
                            d = accumulators::extract::hash<hash_type>(acc);
                        }
                        return (d == _root);
                    }

                    std::size_t leaf_index() const {
                        return _li;
                    }

                    bool operator==(const merkle_proof_impl &rhs) const {
                        return _li == rhs._li && _root == rhs._root && _path == rhs._path;
                    }
                    bool operator!=(const merkle_proof_impl &rhs) const {
                        return !(rhs == *this);
                    }

                    struct path_element_type {
                        path_element_type(value_type x, size_t pos) : _hash(x), _position(pos) {
                        }
                        path_element_type() {
                        }

                        bool operator==(const path_element_type &rhs) const {
                            return _hash == rhs._hash && _position == rhs._position;
                        }
                        bool operator!=(const path_element_type &rhs) const {
                            return !(rhs == *this);
                        }

                        const value_type &hash() const {
                            return _hash;
                        }

                        std::size_t position() const {
                            return _position;
                        }

                        value_type _hash;
                        std::size_t _position;

                        template<typename, typename>
                        friend class nil::crypto3::marshalling::types::merkle_proof_marshalling;
                    };

                    typedef std::array<path_element_type, Arity - 1> layer_type;
                    typedef std::vector<layer_type> path_type;

                    const value_type &root() const {
                        return _root;
                    }

                    const path_type &path() const {
                        return _path;
                    }

                private:
                    std::size_t _li;
                    value_type _root;
                    path_type _path;

                    template<typename, typename, std::size_t>
                    friend class nil::crypto3::zk::components::merkle_proof;

                    template<typename, typename>
                    friend class nil::crypto3::marshalling::types::merkle_proof_marshalling;
                };
            }    // namespace detail

            template<typename T, std::size_t Arity>
            using merkle_proof =
                typename std::conditional<nil::crypto3::detail::is_hash<T>::value,
                                          detail::merkle_proof_impl<detail::merkle_tree_node<T>, Arity>,
                                          detail::merkle_proof_impl<T, Arity>>::type;
        }    // namespace containers
    }        // namespace crypto3
}    // namespace nil

#endif


#include <cctbx/boost_python/flex_fwd.h>
#include <boost/python/module.hpp>
#include <boost/python/class.hpp>
#include <boost/python/def.hpp>
#include <boost/python/args.hpp>
#include <mmtbx/f_model/f_model.h>
#include <scitbx/array_family/boost_python/shared_wrapper.h>
#include <boost/python/return_value_policy.hpp>
#include <boost/python/return_by_value.hpp>

namespace mmtbx { namespace f_model {
namespace {

  boost::python::tuple
  getinitargs(core<> const& self)
  {
    return boost::python::make_tuple(self.f_calc,
                                     self.f_mask(),
                                     self.k_sols(),
                                     self.b_sols(),
                                     self.f_part1,
                                     self.f_part2,
                                     self.u_star,
                                     self.hkl,
                                     self.f_model,
                                     self.f_bulk,
                                     self.k_anisotropic,
                                     self.ss,
                                     self.f_model_no_aniso_scale);
  }

  void init_module()
  {
    using namespace boost::python;
    using boost::python::arg;

    boost::python::to_python_converter<
      af::shared< af::shared<std::complex<double> > > ,
      scitbx::boost_python::container_conversions::to_tuple<
        af::shared< af::shared<std::complex<double> > >   > >();

    scitbx::boost_python::container_conversions::from_python_sequence<
      af::shared< af::shared<std::complex<double> > >,
      scitbx::boost_python::container_conversions::fixed_capacity_policy>();

    typedef return_value_policy<return_by_value> rbv;
    class_<core<> >("core")
      .def(init<
           af::shared<std::complex<double> >      const&,
           af::shared< af::shared<std::complex<double> > > const&,
           af::shared< double>       const&,
           af::shared< double >       const&,
           af::shared<std::complex<double> >      const&,
           af::shared<std::complex<double> >      const&,
           scitbx::sym_mat3<double>               const&,
           af::shared<cctbx::miller::index<> >    const&,
           af::shared<double>                     const& >(
                                                        (arg("f_calc"),
                                                         arg("shell_f_masks"),
                                                         arg("k_sols"),
                                                         arg("b_sols"),
                                                         arg("f_part1"),
                                                         arg("f_part2"),
                                                         arg("u_star"),
                                                         arg("hkl"),
                                                         arg("ss"))))
      .def(init<
           af::shared<std::complex<double> > const&,
           af::shared<std::complex<double> > const&,
           af::shared<double> const&,
           af::shared<double> const&,
           af::shared<double> const&,
           af::shared<std::complex<double> > const&,
           af::shared<std::complex<double> > const& >(
                                         (arg("f_calc"),
                                          arg("f_mask"),
                                          arg("k_isotropic"),
                                          arg("k_anisotropic"),
                                          arg("k_mask"),
                                          arg("f_part1"),
                                          arg("f_part2"))))
      .add_property("f_calc",        make_getter(&core<>::f_calc,       rbv()))
      .add_property("f_part1",       make_getter(&core<>::f_part1,      rbv()))
      .add_property("f_part2",       make_getter(&core<>::f_part2,      rbv()))
      .add_property("u_star",        make_getter(&core<>::u_star,       rbv()))
      .add_property("hkl",           make_getter(&core<>::hkl,          rbv()))
      .add_property("uc",            make_getter(&core<>::uc,           rbv()))
      .add_property("f_model",       make_getter(&core<>::f_model,      rbv()))
      .add_property("f_bulk",        make_getter(&core<>::f_bulk,       rbv()))
      .add_property("k_anisotropic", make_getter(&core<>::k_anisotropic,rbv()))
      .add_property("ss",            make_getter(&core<>::ss,           rbv()))
      .add_property("f_model_no_aniso_scale", make_getter(&core<>::f_model_no_aniso_scale, rbv()))
      .def("n_shells", &core<>::n_shells)
      .def("k_sol", &core<>::k_sol)
      .def("k_sols", &core<>::k_sols)
      .def("b_sols", &core<>::b_sols)
      .def("b_sols", &core<>::b_sols)
      .def("f_mask", &core<>::f_mask)
      .def("k_mask", &core<>::k_mask_one)
      .def("shell_f_mask", &core<>::shell_f_mask)   // XXX re-name to f_mask
      .def("shell_f_masks", &core<>::shell_f_masks) // XXX re-name to f_masks
      .enable_pickling()
      .def("__getinitargs__", getinitargs)
    ;

    typedef return_value_policy<return_by_value> rbv;
    class_<data<> >("data")
      .def(init<
           af::shared<std::complex<double> > const&,
           af::shared<std::complex<double> > const&,
           af::shared<double> const&,
           af::shared<double> const&,
           af::shared<double> const&,
           af::shared<std::complex<double> > const&,
           af::shared<std::complex<double> > const& >(
                                         (arg("f_calc"),
                                          arg("f_bulk"),
                                          arg("k_isotropic_exp"),
                                          arg("k_isotropic"),
                                          arg("k_anisotropic"),
                                          arg("f_part1"),
                                          arg("f_part2"))))
      .add_property("f_calc",        make_getter(&data<>::f_calc,       rbv()))
      .add_property("f_bulk",        make_getter(&data<>::f_bulk,       rbv()))
      .add_property("f_part1",       make_getter(&data<>::f_part1,      rbv()))
      .add_property("f_part2",       make_getter(&data<>::f_part2,      rbv()))
      .add_property("f_model",       make_getter(&data<>::f_model,      rbv()))
      .add_property("k_anisotropic", make_getter(&data<>::k_anisotropic,rbv()))
      .add_property("f_model_no_aniso_scale", make_getter(&data<>::f_model_no_aniso_scale, rbv()))
      .enable_pickling()
      .def("__getinitargs__", getinitargs)
    ;

    def("k_anisotropic",
      (af::shared<double>(*)
        (af::const_ref<cctbx::miller::index<> > const&,
         scitbx::sym_mat3<double> const&)) k_anisotropic);
   ;

   def("k_anisotropic",
      (af::shared<double>(*)
        (af::const_ref<cctbx::miller::index<> > const&,
         af::shared<double> const&,
         cctbx::uctbx::unit_cell const&)) k_anisotropic);
   ;

   def("k_mask",
      (af::shared<double>(*)
        (af::const_ref<double> const&,
         double const&,
         double const&)) k_mask);
   ;

  }

} // namespace <anonymous>
}} // namespace mmtbx::f_model

BOOST_PYTHON_MODULE(mmtbx_f_model_ext)
{
  mmtbx::f_model::init_module();
}


#include "conditionaltest_serial.h"
#include "conditionaltest_workblock.h"
#include "conditionaltest_resultblock.h"
#include "ccmatrix_pair.h"
#include "correlationmatrix_pair.h"
#include "correlationmatrix.h"
#include "correlationmatrix_pair.h"
#include "expressionmatrix_gene.h"
#include <ace/core/elog.h>
#include <gsl/gsl_fit.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_cdf.h>
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_math.h>
#include <boost/math/distributions/non_central_f.hpp>
#include <math.h>
#include <iostream>

/*!
 * Create a serial object.
 */
ConditionalTest::Serial::Serial(ConditionalTest* parent) :
    EAbstractAnalyticSerial(parent),
    _base(parent)
{
    EDEBUG_FUNC(this,parent);
}

/*!
 * Perform the desired work on a result block.
 *
 * @param block The work block you are working on.
 *
 * @return The populated result block.
 */
std::unique_ptr<EAbstractAnalyticBlock> ConditionalTest::Serial::execute(const EAbstractAnalyticBlock* block)
{
    EDEBUG_FUNC(this, block);

    // crate the work and result blocks
    const WorkBlock* workBlock {block->cast<WorkBlock>()};
    ResultBlock* resultBlock {new ResultBlock(workBlock->index())};

    // create iterators for the CCM and CMX data objects
    CCMatrix::Pair ccmPair(_base->_ccm);
    CorrelationMatrix::Pair cmxPair(_base->_cmx);

    // read the first pair of the work block
    Pairwise::Index index(workBlock->start());

    ccmPair.read(index);
    cmxPair.read(index);

    // iterate through each pair in the work block
    for ( qint64 wbIndex = 0; wbIndex < workBlock->size(); ++wbIndex )
    {

        // print warning if iterator indices do not match
        if ( ccmPair.index() != cmxPair.index() )
        {
            QString source = _base->_cmx->geneNames().at(cmxPair.index().getX()).toString();
            QString target = _base->_cmx->geneNames().at(cmxPair.index().getY()).toString();
            qInfo() << "warning: ccm and cmx files are out of sync at cmx coordinate:"
                    << source << "," << target << " ("
                    << cmxPair.index().getX() << "," << cmxPair.index().getY() <<").";
        }

        // initialize set of p values and r2 values for each cluster
        QVector<QVector<double>> pValues(ccmPair.clusterSize());
        QVector<QVector<double>> r2(ccmPair.clusterSize());

        // for each cluster in the pair, run the binomial or linear regression tests
        for ( qint32 clusterIndex = 0; clusterIndex < ccmPair.clusterSize(); clusterIndex++ )
        {

            // resize for room for each test.
            pValues[clusterIndex].resize(_base->_numTests);
            r2[clusterIndex].resize(_base->_numTests);

            // Iterate through all of the tests.
            int test_index = 0;
            for ( qint32 featureIndex = 0; featureIndex < _base->_features.size(); featureIndex++ )
            {
                // Get the column data from the annotation matrix for this feature.
                int num_samples = _base->_data.at(featureIndex).size();
                QVector<QString> amx_column(num_samples);
                for ( int j = 0; j < num_samples; j++ )
                {
                    amx_column[j] = _base->_data.at(featureIndex).at(j).toString();
                }

                if ( _base->_testType.at(featureIndex) == QUANTITATIVE ||
                     _base->_testType.at(featureIndex) == ORDINAL )
                {
                    // For linear regresssion we need a variable that will hold the
                    // pvalue and the r2 value.
                    QVector<double> results(2);
                    regression(amx_column, ccmPair, clusterIndex, results);
                    pValues[clusterIndex][test_index] = results.at(0);
                    r2[clusterIndex][test_index] = results.at(1);
                    test_index++;
                }
                else if ( _base->_testType.at(featureIndex) == CATEGORICAL )
                {
                    // Loop through each label (category) of the feature.
                    for ( qint32 labelIndex = 1; labelIndex < _base->_features.at(featureIndex).size(); labelIndex++ )
                    {
                        double pval;
                        //hypergeom(amx_column, ccmPair, clusterIndex, featureIndex, labelIndex, pval);
                        test_proportions(amx_column, ccmPair, clusterIndex, featureIndex, labelIndex, pval);
                        pValues[clusterIndex][test_index] = pval;
                        r2[clusterIndex][test_index] = qQNaN();
                        test_index++;
                    }
                }
            }
        }

        // append pair to result block
        resultBlock->append(Pair {
            cmxPair.index(),
            pValues,
            r2
        });

        // read the next pair
        cmxPair.readNext();
        ccmPair.read(cmxPair.index());
    }

    return std::unique_ptr<EAbstractAnalyticBlock>(resultBlock);
}

/*!
 * Check to see if a matrix is empty.
 *
 * @param vector The matrix you want to check.
 *
 * @return True if the matrix is empty, false otherwise.
 */
bool ConditionalTest::Serial::isEmpty(QVector<QVector<double>>& matrix)
{
    EDEBUG_FUNC(this, &matrix);

    int index = 0;

    while ( index < matrix.size() )
    {
        if ( !matrix.at(index++).isEmpty() )
        {
            return false;
        }
    }

    return true;
}

/*!
 * The hypergeometric test
 *
 */
void ConditionalTest::Serial::hypergeom(
    const QVector<QString>& amx_column,
    const CCMatrix::Pair& ccmPair,
    int clusterIndex,
    int featureIndex,
    int labelIndex,
    double &result)
{
    EDEBUG_FUNC(this);

    // Get the test name and the test type.
    QString test_label = _base->_features.at(featureIndex).at(labelIndex);
    TestType test_type = _base->_testType.at(featureIndex);

    // We need to get the data for the feature being tested and
    // count how many entries there are for the label in the
    // annotation data.
    int num_samples = _base->_data.at(featureIndex).size();
    int cluster_size = 0;
    int total_label_count = 0;
    int labels_in_cluster = 0;
    for ( int j = 0; j < num_samples; j++ )
    {

        // if data is the same as the test label add one to the catagory counter
        if ( test_type == _base->CATEGORICAL && amx_column[j] == test_label )
        {
            total_label_count++;
        }
        if ( ccmPair.at(clusterIndex, j) == 1 )
        {
            cluster_size++;
            if ( test_type == _base->CATEGORICAL && amx_column.at(j) == test_label )
            {
                labels_in_cluster++;
            }
        }
    }

    // If there are no matching labels in this cluster then return. This
    // could happen if the annotation matrix has all NAs for the cluster.
    if ( labels_in_cluster == 0 )
    {
        result = qQNaN();
        return;
    }

    // We use the hypergeometric distribution because the samples are
    // selected from the population for membership in the cluster without
    // replacement.

    // If a population contains n_1 elements of “type 1” and n_2 elements of
    // “type 2” then the hypergeometric distribution gives the probability
    // of obtaining k elements of “type 1” in t samples from the population.

    int sampleSize =  _base->_emx->sampleSize();

    // Population contains n1 elements of Type 1.
    int n1 = total_label_count;
    // Population contains n2 elements of Type 2.
    int n2 = sampleSize - total_label_count;
    // k elements of Type 1 were selected.
    int k = labels_in_cluster;
    // t total elements were selected.
    int t = cluster_size;

    // If our dataset is large, the power to detect the effect
    // size increases, resulting in potentially insignificant
    // proportions having signficant p-values. Using Cohen's H
    // to set a large-effect size (e.g. 0.8) with a sig.level of
    // 0.001 and a power of 0.95 we need at least 31 samples.
    // So, we'll perform a bootstrap resampling of our data
    // to calculate an average proportion of 31 samples
    int bs_t = 31;
    if ( t > bs_t )
    {
        // Initialize the uniform random number generator.
        const gsl_rng_type * T;
        gsl_rng * r;
        gsl_rng_env_setup();
        T = gsl_rng_default;
        r = gsl_rng_alloc(T);

        // Holds the bootstrap average proportion.
        int jkap = 0;

        // To perform the bootstrap resampling we will
        // perform 30 iterations (central limit thereom)
        int in = 30;
        for ( int i = 0; i < in; i++ )
        {
            // Keeps track of the number of successes for each iteration.
            int ns = 0;

            // Generate 31 random indexes between 0 and the size of
            // the sample string.  We will use these numbers to
            // randomly select a sample and if it is a 1 and of the
            // testing category then we consider it a success.
            int indexes[t];
            int chosen[bs_t];
            int jx = 0;
            for ( int j = 0; j < sampleSize; j++ )
            {
                if ( ccmPair.at(clusterIndex, j) == 1 )
                {
                    indexes[jx++] = j;
                }
            }

            // The gsl_ran_sample function randomly chooses samples with replacement from a list.
            gsl_ran_sample(r, chosen, bs_t, indexes, t, sizeof(int));

            // Now count the number of randomly selected items from the cluster.
            for ( int j = 0; j < bs_t; j++ )
            {
                if ( ccmPair.at(clusterIndex, chosen[j]) == 1 && amx_column.at(chosen[j]) == test_label )
                {
                    ns = ns + 1;
                }
            }

            jkap += ns;
        }

        // Calculate the average proportion from all iterations
        // and free the random number struct.
        jkap = jkap/in;
        gsl_rng_free(r);

        // Since we want to test P(X >= jkap) not P(X > jkap) we should
        // subtract 1 from jkap but we shouldn't go below zero.
        jkap = jkap - 1;
        if (jkap < 0)
        {
            jkap = 0;
        }

        // The gsl_cdf_hypergeometric_Q function uses the upper-tail of the CDF.
        result = gsl_cdf_hypergeometric_Q(jkap, n1, n2, bs_t);
        return;
    }

    // Since we want to test P(X >= k) not P(X > k) we should
    // subtract 1 from k but we shouldn't go below zero.
    k = k - 1;
    if (k < 0)
    {
        k = 0;
    }

    // The gsl_cdf_hypergeometric_Q function uses the upper-tail of the CDF.
    result = gsl_cdf_hypergeometric_Q(k, n1, n2, t);
    return;
}

/*!
 * The z test of proprotions
 *
 */
void ConditionalTest::Serial::test_proportions (
    const QVector<QString>& amx_column,
    const CCMatrix::Pair& ccmPair,
    int clusterIndex,
    int featureIndex,
    int labelIndex,
    double &result)
{
    EDEBUG_FUNC(this);

    // Get the test name and the test type.
    QString test_label = _base->_features.at(featureIndex).at(labelIndex);
    TestType test_type = _base->_testType.at(featureIndex);

    /* For debugging purposes
    QString source_name = _base->_cmx->geneNames().at(ccmPair.index().getX()).toString();
    QString target_name = _base->_cmx->geneNames().at(ccmPair.index().getY()).toString();

    if (source_name == "pycom14g13710" && target_name == "pycom02g20220") {
      std::cout << "Found it";
    }*/

    // We need to get the data for the feature being tested and
    // count how many entries there are for the label in the
    // annotation data.
    int num_samples = _base->_data.at(featureIndex).size();
    int cluster_size = 0;
    int total_label_count = 0;
    int total_non_missing = 0;
    int labels_in_cluster = 0;
    int labels_out_cluster = 0;
    int jx = 0;
    for ( int j = 0; j < num_samples; j++ )
    {

        // if data is the same as the test label add one to the catagory counter
        if ( test_type == _base->CATEGORICAL && amx_column[j] == test_label )
        {
            total_label_count++;
        }
        if ( ccmPair.at(clusterIndex, j) == 1 )
        {
            cluster_size++;
            if ( test_type == _base->CATEGORICAL && amx_column.at(j) == test_label )
            {
                labels_in_cluster++;
            }
        }
        if ( ccmPair.at(clusterIndex, j) != 1 && ccmPair.at(clusterIndex, j) != 9)
        {
            if ( test_type == _base->CATEGORICAL && amx_column.at(j) == test_label )
            {
                labels_out_cluster++;
            }
        }
        if ( ccmPair.at(clusterIndex, j) != 9 )
        {
          total_non_missing++;
        }
    }

    // If there are no matching labels in this cluster then return. This
    // could happen if the annotation matrix has all NAs for the cluster.
    if ( labels_in_cluster == 0 )
    {
        result = qQNaN();
        return;
    }

    int sampleSize =  _base->_emx->sampleSize();

    // Before bootstraping get the list of indexes that belong to this cluster.
    // We will randomly select from these indexes.
    int cluster_indexes[cluster_size];
    for ( int j = 0; j < sampleSize; j++ )
    {
        if ( ccmPair.at(clusterIndex, j) == 1 )
        {
            cluster_indexes[jx++] = j;
        }
    }

    // Initialize the uniform random number generator.
    const gsl_rng_type * T;
    gsl_rng * r;
    gsl_rng_env_setup();
    T = gsl_rng_default;
    r = gsl_rng_alloc(T);

    // Perform 30 iterations of the boostrap sampling.
    int num_i = 30;
    // Select 30 items for each iteration
    int bs_t = 30;
    // Holds the bootstrap average proportion.
    double avg_prop = 0.0;
    for ( int i = 0; i < num_i; i++ )
    {
        // Keeps track of the number of successes for each iteration.
        int ns = 0;

        // The gsl_ran_sample function randomly chooses samples with replacement from a list.
        int chosen[bs_t];
        gsl_ran_sample(r, chosen, bs_t, cluster_indexes, cluster_size, sizeof(int));

        // Now count the number of randomly selected items from the cluster.
        for ( int j = 0; j < bs_t; j++ )
        {
            QString slabel = amx_column.at(chosen[j]);
            if ( ccmPair.at(clusterIndex, chosen[j]) == 1 && amx_column.at(chosen[j]) == test_label )
            {
                ns = ns + 1;
            }
        }

        avg_prop += static_cast<double>(ns) / static_cast<double>(bs_t);
    }

    // Calculate the average proportion from all iterations
    // and free the random number struct.
    avg_prop = avg_prop / num_i;
    gsl_rng_free(r);

    // First test.
    // The null hypothesis for this test is that the cluster is
    // comprised of equal or less of the desired label. This would
    // happen if all labels have equal chance, or there is a bias
    // against the label.
    // Ho: avg_prop <= 0.5. The value 0.5 is chosen with the asumption
    //     that there is equal or less proportions of the desired label.
    // Ha: avg_prop > 0.5
    double exp_prop = 0.5;
    double z = (avg_prop - exp_prop) / sqrt((exp_prop * (1 - exp_prop)) / bs_t);
    double pval1 = 1 - gsl_cdf_gaussian_P(z, 1);

    // Second, test.
    // The null hypothesis is that the proportion of labels in the cluster
    // is equal or less than the proportion out of the cluster.  This would
    // happen if all the labels have equal chance of being in the cluster or
    // if there is bias against the label.
    double obs_prop = (avg_prop * cluster_size) / static_cast<double>(total_label_count);
    obs_prop = std::min(1.0, obs_prop);
    exp_prop = static_cast<double>(total_label_count) / static_cast<double>(num_samples);
    z = (obs_prop - exp_prop) / sqrt((exp_prop * (1 - exp_prop)) / bs_t);
    double pval2 = 1 - gsl_cdf_gaussian_P(z, 1);

    // Return the maximum p-value
    result = std::max(pval1, pval2);
}

/*!
 * Performs the regression test for quantitative data.
 */
void ConditionalTest::Serial::regression(
    const QVector<QString>& amx_column,
    const CCMatrix::Pair& ccmPair,
    int clusterIndex,
    QVector<double>& results)
{
    EDEBUG_FUNC(this, &amxInfo, &ccmPair, clusterIndex);

    // Temp containers.
    QVector<double> labelInfo;

    // Regression model containers.
    double chisq;
    gsl_matrix *X, *cov;
    gsl_vector *Y, *C;

    // Before allocating memory for the linear regression let's remove any
    // values in the expression matrix that have missing values. This will
    // tell us the full size we need to reserve.
    int test_cluster_size = 0;
    for ( int i = 0; i < _base->_emx->sampleSize(); i++ )
    {
        // If the sample label matches with the given label.
        if ( ccmPair.at(clusterIndex, i) == 1 )
        {
            if (QString::compare(amx_column.at(i), _base->_missing) == 0) {
                continue;
            }
            test_cluster_size++;
        }
    }

    // Don't do regression analysis if there are fewer than 5 samples.
    // This can occur if the annotation matrix has missing values
    // for samples in the cluster. The Degrees of Freedom for the error
    // is test_cluster_size - 4 so we need at least 5 samples. Odds
    // are these tests will get filterd out anyway for low power.
    if ( test_cluster_size < 5 )
    {
        results[0] = qQNaN();
        results[1] = qQNaN();
        return;
    }

    // Allocate a matrix to hold the predictior variables, in this case the gene
    // Expression data.
    X = gsl_matrix_alloc(test_cluster_size, 4);

    // Allocate a vector to hold observation data, in this case the data
    // corrosponding to the features.
    Y = gsl_vector_alloc(test_cluster_size);

    // Allocate a vector and matrix for the slope info.
    C = gsl_vector_alloc(4);
    cov = gsl_matrix_alloc(4, 4);

    // Read in the gene pairs expression information.
    ExpressionMatrix::Gene geneX(_base->_emx);
    ExpressionMatrix::Gene geneY(_base->_emx);
    geneX.read(ccmPair.index().getX());
    geneY.read(ccmPair.index().getY());

    // Useful for debugging purposes.
    // QString g1Name = geneX.toString();
    // QString g2Name = geneY.toString();

    // Look through all the samples in the mask.
    for ( int i = 0, j = 0; i < _base->_emx->sampleSize(); i++ )
    {
        // If the sample label matches with the given label.
        if ( ccmPair.at(clusterIndex, i) == 1 )
        {
            // Skip samples with a missing value in the annotation matrix
            if (QString::compare(amx_column.at(i), _base->_missing) == 0) {
                continue;
            }

            // Add emx data as the predictors but only for samples in the cluster. We
            // add a 1 for the intercept, gene1, gene2 and the interaction: gene1*gene2.
            // The interaction term handles the case where the relationship is dependent
            // on the values of both genes.
            double g1 = static_cast<double>(geneX.at(i));
            double g2 = static_cast<double>(geneY.at(i));

            gsl_matrix_set(X, j, 0, static_cast<double>(1)); // for the intercept
            gsl_matrix_set(X, j, 1, g1);
            gsl_matrix_set(X, j, 2, g2);
            gsl_matrix_set(X, j, 3, g1*g2);
            gsl_vector_set(Y, j, amx_column.at(i).toFloat());

            j++;
        }
    }

    // Create the workspace for the gnu scientific library to work in.
    gsl_multifit_linear_workspace * work = gsl_multifit_linear_alloc (X->size1, X->size2);

    // Regression calculation.
    gsl_multifit_linear(X, Y, C, cov, &chisq, work);

    // Calculate some quantities we'll need for R^2 and p-value
    // calculations.
    // Sum of Squares Total
    double SST = gsl_stats_tss(Y->data, Y->stride, Y->size);
    // Sum of Squares of the Error (residuals)
    double SSE = chisq;

    // Sum of Squares for the model.
    double SSM = SST - SSE;

    // Degrees of Freedom for the Model. This is p-1 where
    // p is the number of parameters in the model.
    double DFM = 3;
    // Degrees of Freedom for the Error. This is n-p where
    // n is the number of samples.
    double DFE = test_cluster_size - 4;

    // Mean of Squares for the Model.
    double MSM = SSM/DFM;
    // Mean of Squares for Error.
    double MSE = SSE/DFE;

    // Calculate R^2 and R^2-adjusted
    double R2 = 1 - (SSE / SST);
    //double R2adj = 1.0 - ((double) test_cluster_size - 1) / DFE * (1 - R2);

    // Calculate the p-value. We will do this using the F-test.
    double Fstat = MSM/MSE;
    double pValue = 1 - gsl_cdf_fdist_P(Fstat, DFM, DFE);

    // TODO: we should check the assumptions of the linear regression and
    // not return if the assumptions are not met.

    // TODO: it would be nice to return a rate of change of the conditioal mean.

    // TODO: it would be nice to return the r2 value along with the p-value.

    // Four scenarios:
    // 1) low R-square and low p-value (p-value <= 0.05).  Model doesn't explain
    //    the variation but it does follow the trend or regression line well.
    // 2) low R-square and high p-value (p-value > 0.05).  Model doesn't explai
    //    the variation and it doesn't follow the trend line.
    // 3) high R-square and low p-value.  Model explains the variance it
    //    follows the trend line.
    // 4) high R-square and high p-value. Model explains the variance well but
    //    it does not follow the trend line very well.
    // In summary, low p-values still indicate a real relationship between the
    // predictors and the observed values.

    // Free all of the data.
    gsl_matrix_free(X);
    gsl_vector_free(Y);
    gsl_matrix_free(cov);
    gsl_vector_free(C);
    gsl_multifit_linear_free(work);

    // Set the results array
    if ( qIsNaN(pValue) )
    {
        results[0] = qQNaN();
        results[1] = qQNaN();
    }
    else {
        double power = pwr_f2_test(3, test_cluster_size - 4, R2, _base->_powerThresholdAlpha);
        if ( power >= _base->_powerThresholdPower ) {
            results[0] = pValue;
            results[1] = R2;
        }
        else {
            results[0] = std::numeric_limits<double>::quiet_NaN();
            results[1] = std::numeric_limits<double>::quiet_NaN();
        }
    }
}

/**
 * @brief ConditionalTest::Serial::pwr_f2_test
 * @param df1 degrees of freedom of the numerator: the number of coefficients in
 *        the model minus the intercept.
 * @param df2 degrees of freedom of the denominator: the number of error degrees
 *        of freedom which is n - df1 - 1
 * @param f2 effect size
 * @param sig_level signficance level
 * @return
 */
double ConditionalTest::Serial::pwr_f2_test(int df1, int df2, double f2, double sig_level)
{
    double lambda = f2 * (df1 + df2 + 1);
    double p = gsl_cdf_fdist_Qinv(sig_level, df1, df2);
    boost::math::non_central_f_distribution<> ncfd(df1, df2, lambda);
    double power = 1 - boost::math::cdf(ncfd, p);
    return power;
}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_IEEE_FUNCTIONS_SIMD_SSE_SSE2_ILOGB_HPP_INCLUDED
#define BOOST_SIMD_IEEE_FUNCTIONS_SIMD_SSE_SSE2_ILOGB_HPP_INCLUDED
#ifdef BOOST_SIMD_HAS_SSE2_SUPPORT
#include <boost/simd/ieee/functions/ilogb.hpp>
#include <boost/simd/ieee/functions/simd/common/ilogb.hpp>
#include <boost/simd/include/constants/zero.hpp>
#include <boost/simd/include/functions/simd/tofloat.hpp>
#include <boost/simd/include/functions/simd/seladd.hpp>
#include <boost/simd/include/functions/simd/is_gtz.hpp>
#include <boost/simd/include/functions/simd/is_nez.hpp>
#include <boost/simd/include/functions/simd/if_else.hpp>
#include <boost/simd/include/functions/simd/predecessor.hpp>
#include <boost/simd/include/functions/simd/dist.hpp>
#define MKN(N) simd::bitwise_cast<vtype##N>

namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::ilogb_, boost::simd::tag::sse2_,
                        (A0),
                        ((simd_<int_<A0>,boost::simd::tag::sse_>))
                       )
  {
    typedef typename dispatch::meta::as_integer<A0>::type  result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(1)
    {
      typedef typename dispatch::meta::as_integer<A0,unsigned>::type vtype;
      static const A0 z = Zero<A0>();
      vtype tmp = ilogb(simd::bitwise_cast<vtype>(a0));
      return seladd(is_gtz(a0), z, simd::bitwise_cast<A0>(tmp));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::ilogb_, boost::simd::tag::sse2_,
                        (A0),
                        ((simd_<uint8_<A0>,boost::simd::tag::sse_>))
                       )
  {
    typedef typename dispatch::meta::as_integer<A0>::type  result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(1)
    {
      typedef A0 vtype8;
      static const vtype8& o = One<vtype8>();
      static const vtype8 mask =  boost::simd::integral_constant<vtype8, 0x7f>();
      typedef typename A0::extension_type cat;
      typedef boost::simd::native<boost::uint16_t, cat> type;
      vtype8 n = a0;
      vtype8 i = o;
      n = b_and(MKN(8)(_mm_srli_epi16(n, 1)), mask);
      i = seladd(is_nez(n), i, o);
      n = b_and(MKN(8)(_mm_srli_epi16(n, 1)), mask);
      i = seladd(is_nez(n), i, o);
      n = b_and(MKN(8)(_mm_srli_epi16(n, 1)), mask);
      i = seladd(is_nez(n), i, o);
      n = b_and(MKN(8)(_mm_srli_epi16(n, 1)), mask);
      i = seladd(is_nez(n), i, o);
      n = b_and(MKN(8)(_mm_srli_epi16(n, 1)), mask);
      i = seladd(is_nez(n), i, o);
      n = b_and(MKN(8)(_mm_srli_epi16(n, 1)), mask);
      i = seladd(is_nez(n), i, o);
      n = b_and(MKN(8)(_mm_srli_epi16(n, 1)), mask);
      i = seladd(is_nez(n), i, o);
      return i-o;
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::ilogb_, boost::simd::tag::sse2_,
                        (A0),
                        ((simd_<uint32_<A0>,boost::simd::tag::sse_>))
                       )
  {
    typedef typename dispatch::meta::as_integer<A0>::type  result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(1)
    {
      typedef typename A0::extension_type cat;
      typedef A0 vtype32;
      typedef simd::native<boost::uint8_t, cat>  vtype8;
      static const vtype8& z = Zero<vtype8>();
      static const vtype8& o = One<vtype8>();
      static const vtype32 mask = boost::simd::integral_constant<vtype32, 0x7f7f7f7f>();
      vtype8 i = if_else(is_eqz(MKN(8)(a0)), z, o);
      vtype8 n = MKN(8)(b_and(MKN(32)(_mm_srli_epi16(a0, 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(32)(_mm_srli_epi16(MKN(32)(n), 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(32)(_mm_srli_epi16(MKN(32)(n), 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(32)(_mm_srli_epi16(MKN(32)(n), 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(32)(_mm_srli_epi16(MKN(32)(n), 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(32)(_mm_srli_epi16(MKN(32)(n), 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(32)(_mm_srli_epi16(MKN(32)(n), 1)), mask));
      i = seladd(is_nez(n), i, o);
      static const vtype32 mask1 = boost::simd::integral_constant<vtype32,0x000000ff>();
      static const vtype32 mask2 = boost::simd::integral_constant<vtype32,0x0000ff00>();
      static const vtype32 mask3 = boost::simd::integral_constant<vtype32,0x00ff0000>();
      static const vtype32 mask4 = boost::simd::integral_constant<vtype32,0xff000000>();
      static const vtype32 huit  = boost::simd::integral_constant<vtype32,8>();
      static const vtype32 seize = boost::simd::integral_constant<vtype32,16u>();
      static const vtype32 v4    = boost::simd::integral_constant<vtype32,24u>();
      vtype32 yy = b_and(MKN(32)(i), mask4);
      vtype32 zz = b_and(MKN(32)(i), mask3);
      vtype32 tt = b_and(MKN(32)(i), mask2);
      vtype32 uu = b_and(MKN(32)(i), mask1);
      vtype32 xx = if_else(is_nez(yy), MKN(32)(_mm_srli_epi32(yy,24))+v4,
                 if_else(is_nez(zz), MKN(32)(_mm_srli_epi32(zz,16))+seize,
                   if_else(is_nez(tt), MKN(32)(_mm_srli_epi32(tt,8))+huit,
                       uu)));
      return MKN(32)(is_nez(MKN(8)(xx))&(MKN(8)(xx)-o));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::ilogb_, boost::simd::tag::sse2_,
                        (A0),
                        ((simd_<uint64_<A0>,boost::simd::tag::sse_>))
                       )
  {
    typedef typename dispatch::meta::as_integer<A0>::type  result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(1)
    {
      typedef typename A0::extension_type cat;
      typedef simd::native<boost::uint64_t, cat> vtype64;
      return simd::bitwise_cast<vtype64>(exponent(tofloat(a0)));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::ilogb_, boost::simd::tag::sse2_,
                        (A0),
                        ((simd_<uint16_<A0>,boost::simd::tag::sse_>))
                       )
  {
    typedef typename dispatch::meta::as_integer<A0>::type  result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(1)
    {
      typedef A0 vtype16;
      typedef typename A0::extension_type cat;
      typedef simd::native<boost::uint8_t, cat>  vtype8;
      static const vtype8& z = Zero<vtype8>();
      static const vtype8& o = One<vtype8>();
      static const vtype16 mask =  boost::simd::integral_constant<vtype16, 0x7f7f > ();
      vtype8 i = if_else(is_eqz(MKN(8)(a0)), z, o);
      vtype8 n = MKN(8)(b_and(MKN(16)(_mm_srli_epi16(a0, 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(16)(_mm_srli_epi16(n, 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(16)(_mm_srli_epi16(n, 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(16)(_mm_srli_epi16(n, 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(16)(_mm_srli_epi16(n, 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(16)(_mm_srli_epi16(n, 1)), mask));
      i = seladd(is_nez(n), i, o);
      n = MKN(8)(b_and(MKN(16)(_mm_srli_epi16(n, 1)), mask));
      i = seladd(is_nez(n), i, o);
      static const vtype16 mask1 = boost::simd::integral_constant<vtype16, 0x00ff>();
      static const vtype16 mask2 = boost::simd::integral_constant<vtype16, 0xff00>();
      static const vtype16 huit  = boost::simd::integral_constant<vtype16, 8u>();
      vtype16 yy = b_and(MKN(16)(i), mask2);
      vtype16 zz = b_and(MKN(16)(i), mask1);
      vtype16 xx = if_else(is_nez(yy), MKN(16)(_mm_srli_epi16(yy, 8))+huit, zz);
      return MKN(16)(is_nez(MKN(8)(xx))&(MKN(8)(xx)-o));
    }
  };
} } }

#endif
#endif


// Copyright 2019-2020 Josh Pieper, jjp@pobox.com.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "moteus/tool/calibrate.h"

#include <algorithm>
#include <cmath>
#include <deque>

#include <fmt/format.h>

#include <Eigen/Core>

#include <boost/algorithm/string.hpp>

#include "mjlib/base/system_error.h"

namespace moteus {
namespace tool {

namespace {

constexpr double kPi = 3.14159265358979323846;

struct Entry {
  int phase = 0;
  int encoder = 0;
};

struct File {
  std::vector<Entry> phase_up;
  std::vector<Entry> phase_down;
};

File ParseFile(const std::vector<std::string>& lines_in) {
  std::deque<std::string> lines;
  std::copy(lines_in.begin(), lines_in.end(), std::back_inserter(lines));
  if (lines.front() != "CAL start") {
    mjlib::base::system_error::throw_if(
        true, "Does not start with magic line: " + lines.front());
  }

  lines.pop_front();

  File result;

  while (!lines.empty()) {
    const auto line = lines.front();
    lines.pop_front();

    if (line == "CAL done") {
      return result;
    }

    std::vector<std::string> fields;
    boost::split(fields, line, boost::is_any_of(" "));

    if (fields.size() != 3) {
      mjlib::base::system_error::throw_if(
          true, "malformed line: " + line);
    }

    const auto phase = std::stoi(fields[1]);
    const auto encoder = std::stoi(fields[2]);
    const auto entry = Entry{phase, encoder};

    if (fields[0] == "1") {
      result.phase_up.push_back(entry);
    } else if (fields[0] == "2") {
      result.phase_down.push_back(entry);
    } else {
      mjlib::base::system_error::throw_if(
          true, "unknown phase code: " + line);
    }
  }

  mjlib::base::system_error::throw_if(
      true, "Does not end with magic line");
  return {};
}

int WrapInt16(int value) {
  if (value > 32767) {
    return value - 65536;
  } else if (value < -32768) {
    return 65536 + value;
  }
  return value;
}

double WrapNegPiToPi(double value) {
  if (value >= -kPi && value <= kPi) { return value; }
  if (value > 0.0) {
    return std::fmod(value + kPi, 2 * kPi) - kPi;
  } else {
    return std::fmod(value - kPi, 2 * kPi) + kPi;
  }
}

Eigen::VectorXd Unwrap(const Eigen::VectorXd& value) {
  Eigen::VectorXd result(value.size());
  for (int i = 0; i < value.size(); i++) {
    if (i == 0) {
      result(i) = value(i);
    } else {
      result(i) = (result(i - 1) + WrapNegPiToPi(value(i) - result(i - 1)));
    }
  }

  return result;
}

Eigen::VectorXd Linspace(double start, double end, size_t count) {
  Eigen::VectorXd result(count);
  for (size_t i = 0; i < count; i++) {
    result(i) = (
        (i + 1) == count ?
        end :
        i * (end - start) / (count - 1) + start);
  }
  return result;
}

Eigen::VectorXd Range(double start, double end, size_t count) {
  Eigen::VectorXd result(count);
  for (size_t i = 0; i < count; i++) {
    result(i) = (i * (end - start) / count + start);
  }
  return result;
}

Eigen::VectorXd Interpolate(const Eigen::VectorXd& sample_points,
                            const Eigen::VectorXd& x,
                            const Eigen::VectorXd& y) {
  BOOST_ASSERT(x.size() > 1 && x.size() == y.size());

  int xindex = 0;

  Eigen::VectorXd result(sample_points.size());
  for (int i = 0; i < sample_points.size(); i++) {
    const auto point = sample_points[i];
    result(i) = [&]() {
      if (point < x[xindex]) {
        // Whoops?  This can only legimately happen at the very beginning.
        return y[xindex];
      }
      while ((xindex + 2) < x.size() &&
             point >= x[xindex + 1]) {
        xindex++;
      }

      if (point > x[xindex + 1]) {
        // We're past the end.
        return y[xindex + 1];
      }

      // Linearly interpolate.
      const double length = x[xindex + 1] - x[xindex];
      if (length == 0.0) {
        return y[xindex + 1];
      }
      double ratio = (point - x[xindex]) / length;
      return (y[xindex + 1] - y[xindex]) * ratio + y[xindex];
    }();
  }

  return result;
}

Eigen::VectorXd WindowAverage(const Eigen::VectorXd& values, int window_size) {
  Eigen::VectorXd result(values.size());
  for (int i = 0; i < values.size(); i++) {
    int start = std::max<int>(0, i - window_size / 2);
    int end = std::min<int>(values.size(), i + window_size / 2);
    Eigen::VectorXd errs(end - start);
    for (int i = start; i < end; i++) {
      errs[i - start] = WrapNegPiToPi(values[i] - values[start]);
    }
    result[i] = values[start] + errs.mean();
  }

  return result;
}

}

CalibrationResult Calibrate(const std::vector<std::string>& lines) {
  auto file = ParseFile(lines);

  if (file.phase_up.empty() ||
      file.phase_down.empty()) {
    mjlib::base::system_error::throw_if(
        true, "one or more phases was empty");
  }

  double total_delta = 0;
  for (size_t i = 0; i < file.phase_up.size() - 1; i++) {
    total_delta += WrapInt16(file.phase_up[i + 1].encoder -
                             file.phase_up[i].encoder);
  }

  if (std::abs(std::abs(total_delta) - 65536) > 5000) {
    mjlib::base::system_error::throw_if(
        true, "phase_up did not traverse appropriate encoder distance");
  }

  CalibrationResult result;

  // Figure out inversion.
  if (total_delta < 0) {
    result.invert = true;
    for (auto& item : file.phase_up) { item.encoder = 65535 - item.encoder; }
    for (auto& item : file.phase_down) { item.encoder = 65535 - item.encoder; }
    total_delta *= -1;
  }

  // Next, figure out the number of poles.  We compare the total
  // encoder delta to the total phase delta.
  double total_phase = 0;
  for (size_t i = 0; i < file.phase_up.size() - 1; i++) {
    total_phase += WrapInt16(file.phase_up[i + 1].phase -
                             file.phase_up[i].phase);
  }

  const double ratio = total_phase / total_delta;
  const double remainder = std::abs(std::round(ratio) - ratio);

  const double kMaxRemainderError = 0.1;
  if (remainder > kMaxRemainderError) {
    mjlib::base::system_error::throw_if(
        true, fmt::format(
            "encoder not an integral multiple of phase, {} > {}",
            remainder, kMaxRemainderError));
  }

  result.total_phase = total_phase;
  result.total_delta = total_delta;
  result.ratio = ratio;

  result.poles = static_cast<int>(std::round(ratio) * 2);

  // Now we need to figure out the phase offset at select points.  We
  // interpolate and average the phase up and phase down sections.
  std::vector<Entry> phase_up_by_encoder = file.phase_up;
  const auto encoder_sort = [](const auto& lhs, const auto& rhs) {
    return lhs.encoder < rhs.encoder;
  };
  std::sort(phase_up_by_encoder.begin(), phase_up_by_encoder.end(),
            encoder_sort);
  std::vector<Entry> phase_down_by_encoder = file.phase_down;
  std::sort(phase_down_by_encoder.begin(), phase_down_by_encoder.end(),
            encoder_sort);

  double offset = phase_down_by_encoder[0].phase - phase_up_by_encoder[0].phase;
  if (std::abs(offset) > 32767) {
    // We need to shift it so that they start from the same place.
    auto change = static_cast<int>(-65536 * std::round(offset / 65536.0));
    for (auto& item : phase_down_by_encoder) { item.phase += change; }
  }

  Eigen::VectorXd phase_up_encoder(phase_up_by_encoder.size());
  for (size_t i = 0; i < phase_up_by_encoder.size(); i++) {
    phase_up_encoder[i] = phase_up_by_encoder[i].encoder;
  }

  Eigen::VectorXd phase_up_phase(phase_up_by_encoder.size());
  for (size_t i = 0; i < phase_up_by_encoder.size(); i++) {
    phase_up_phase[i] =
        2.0 * kPi / 65536.0 * phase_up_by_encoder[i].phase;
  }
  phase_up_phase = Unwrap(phase_up_phase);

  Eigen::VectorXd phase_down_encoder(phase_down_by_encoder.size());
  for (size_t i = 0; i < phase_down_by_encoder.size(); i++) {
    phase_down_encoder[i] = phase_down_by_encoder[i].encoder;
  }

  Eigen::VectorXd phase_down_phase(phase_down_by_encoder.size());
  for (size_t i = 0; i < phase_down_by_encoder.size(); i++) {
    phase_down_phase[i] =
        2.0 * kPi / 65536.0 * phase_down_by_encoder[i].phase;
  }
  phase_down_phase = Unwrap(phase_down_phase);

  Eigen::VectorXd xpos = Linspace(0, 65535.0, 10000);

  const Eigen::VectorXd pu_interp =
      Interpolate(xpos, phase_up_encoder, phase_up_phase);
  const Eigen::VectorXd pd_interp =
      Interpolate(xpos, phase_down_encoder, phase_down_phase);
  const Eigen::VectorXd avg_interp = 0.5 * (pu_interp + pd_interp);

  const Eigen::VectorXd expected =
      (2.0 * kPi / 65536) * (result.poles / 2) * xpos;

  Eigen::VectorXd err(expected.size());
  for (int i = 0; i < expected.size(); i++) {
    err[i] = WrapNegPiToPi(avg_interp[i] - expected[i]);
  }

  // Make the error seem reasonable, so unwrap if we happen to span
  // the pi boundary.
  if ((err.maxCoeff() - err.minCoeff()) > 1.5 * kPi) {
    for (int i = 0; i < err.size(); i++) {
      err[i] = (err[i] > 0) ? err[i] : (err[i] + 2 * kPi);
    }
  }

  auto avg_window = static_cast<int>(err.size() / result.poles);
  const Eigen::VectorXd avg_err = WindowAverage(err, avg_window);

  const Eigen::VectorXd offset_x = Range(0, 65536, 64);
  const Eigen::VectorXd offsets = Interpolate(offset_x, xpos, avg_err);

  for (int i = 0; i < offsets.size(); i++) {
    result.offset.push_back(offsets[i]);
  }

  return result;
}

}
}


// Boost.Geometry (aka GGL, Generic Geometry Library)
// Unit Test

// Copyright (c) 2010-2012 Barend Gehrels, Amsterdam, the Netherlands.

// This file was modified by Oracle on 2020-2021.
// Modifications copyright (c) 2020-2021, Oracle and/or its affiliates.
// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#include <iostream>
#include <string>

#define BOOST_GEOMETRY_UNIT_TEST_SECTIONALIZE

#include <geometry_test_common.hpp>

#include <boost/geometry/algorithms/detail/sections/sectionalize.hpp>
#include <boost/geometry/algorithms/detail/sections/range_by_section.hpp>
#include <boost/geometry/geometries/geometries.hpp>
#include <boost/geometry/geometries/point_xy.hpp>
#include <boost/geometry/io/wkt/wkt.hpp>
#include <boost/geometry/util/condition.hpp>

// TEMP
#include <boost/geometry/strategies/cartesian.hpp>

template <int DimensionCount, bool Reverse, typename Geometry>
void test_sectionalize(std::string const /*caseid*/, Geometry const& geometry, std::size_t section_count)
{
    typedef typename bg::point_type<Geometry>::type point;
    typedef bg::model::box<point> box;
    typedef bg::sections<box, DimensionCount> sections;

    typedef std::integer_sequence<std::size_t, 0> dim2;

    sections s;
    bg::sectionalize<Reverse, dim2>(geometry, bg::detail::no_rescale_policy(), s);

    BOOST_CHECK_EQUAL(s.size(), section_count);

    typedef typename bg::closeable_view
        <
            typename bg::ring_type<Geometry>::type const,
            bg::closure<Geometry>::value
        >::type cview_type;
    typedef typename bg::reversible_view
        <
            cview_type const,
            Reverse ? bg::iterate_reverse : bg::iterate_forward
        >::type view_type;
    typedef typename boost::range_iterator
        <
            view_type const
        >::type range_iterator;

    BOOST_FOREACH(typename sections::value_type const& sec, s)
    {
        cview_type cview(bg::range_by_section(geometry, sec));
        view_type view(cview);
        range_iterator it1 = boost::begin(view) + sec.begin_index;
        range_iterator it2 = boost::begin(view) + sec.end_index;
        int count = 0;
        for (range_iterator it = it1; it != it2; ++it)
        {
            count++;
        }
        BOOST_CHECK_EQUAL(int(sec.count), count);
    }
}

template <typename Geometry, bool Reverse>
void test_sectionalize(std::string const& caseid, std::string const& wkt,
        std::size_t count1)
{
    Geometry geometry;
    bg::read_wkt(wkt, geometry);
    if ( BOOST_GEOMETRY_CONDITION( bg::closure<Geometry>::value == bg::open ) )
    {
        geometry.outer().resize(geometry.outer().size() - 1);
    }
    //bg::correct(geometry);
    test_sectionalize<1, Reverse>(caseid + "_d1", geometry, count1);
}

template <typename P>
void test_all()
{
    std::string const first = "polygon((2.0 1.3, 2.4 1.7, 2.8 1.8, 3.4 1.2, 3.7 1.6,3.4 2.0, 4.1 3.0, 5.3 2.6, 5.4 1.2, 4.9 0.8, 2.9 0.7,2.0 1.3))";
    test_sectionalize<bg::model::polygon<P>, false>("first", first, 4);

    test_sectionalize<bg::model::polygon<P, false>, true>("first_reverse",
        first, 4);

    test_sectionalize<bg::model::polygon<P, false, true>, false>("first_open",
        first, 4);

    test_sectionalize<bg::model::polygon<P, true, false>, true>("first_open_reverse",
        first, 4);
}

int test_main(int, char* [])
{
    test_all<bg::model::d2::point_xy<double> >();

    return 0;
}