/*
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/** \file
 * \ingroup bke
 */

#include "MEM_guardedalloc.h"

#include <stdlib.h>

#include "BLI_dlrbTree.h"
#include "BLI_listbase.h"
#include "BLI_math.h"

#include "DNA_anim_types.h"
#include "DNA_armature_types.h"
#include "DNA_scene_types.h"

#include "BKE_action.h"
#include "BKE_anim_data.h"
#include "BKE_main.h"
#include "BKE_scene.h"

#include "DEG_depsgraph.h"
#include "DEG_depsgraph_build.h"
#include "DEG_depsgraph_query.h"

#include "GPU_batch.h"
#include "GPU_vertex_buffer.h"

#include "ED_anim_api.h"
#include "ED_keyframes_draw.h"

#include "CLG_log.h"

static CLG_LogRef LOG = {"ed.anim.motion_paths"};

/* Motion path needing to be baked (mpt) */
typedef struct MPathTarget {
  struct MPathTarget *next, *prev;

  bMotionPath *mpath; /* motion path in question */

  DLRBT_Tree keys; /* temp, to know where the keyframes are */

  /* Original (Source Objects) */
  Object *ob;          /* source object */
  bPoseChannel *pchan; /* source posechannel (if applicable) */

  /* "Evaluated" Copies (these come from the background COW copy
   * that provide all the coordinates we want to save off). */
  Object *ob_eval; /* evaluated object */
} MPathTarget;

/* ........ */

/* update scene for current frame */
static void motionpaths_calc_update_scene(struct Depsgraph *depsgraph)
{
  BKE_scene_graph_update_for_newframe(depsgraph);
}

Depsgraph *animviz_depsgraph_build(Main *bmain,
                                   Scene *scene,
                                   ViewLayer *view_layer,
                                   ListBase *targets)
{
  /* Allocate dependency graph. */
  Depsgraph *depsgraph = DEG_graph_new(bmain, scene, view_layer, DAG_EVAL_VIEWPORT);

  /* Make a flat array of IDs for the DEG API. */
  const int num_ids = BLI_listbase_count(targets);
  ID **ids = MEM_malloc_arrayN(sizeof(ID *), num_ids, "animviz IDS");
  int current_id_index = 0;
  for (MPathTarget *mpt = targets->first; mpt != NULL; mpt = mpt->next) {
    ids[current_id_index++] = &mpt->ob->id;
  }

  /* Build graph from all requested IDs. */
  DEG_graph_build_from_ids(depsgraph, ids, num_ids);
  MEM_freeN(ids);

  /* Update once so we can access pointers of evaluated animation data. */
  motionpaths_calc_update_scene(depsgraph);
  return depsgraph;
}

/* get list of motion paths to be baked for the given object
 * - assumes the given list is ready to be used
 */
/* TODO: it would be nice in future to be able to update objects dependent on these bones too? */
void animviz_get_object_motionpaths(Object *ob, ListBase *targets)
{
  MPathTarget *mpt;

  /* object itself first */
  if ((ob->avs.recalc & ANIMVIZ_RECALC_PATHS) && (ob->mpath)) {
    /* new target for object */
    mpt = MEM_callocN(sizeof(MPathTarget), "MPathTarget Ob");
    BLI_addtail(targets, mpt);

    mpt->mpath = ob->mpath;
    mpt->ob = ob;
  }

  /* bones */
  if ((ob->pose) && (ob->pose->avs.recalc & ANIMVIZ_RECALC_PATHS)) {
    bArmature *arm = ob->data;
    bPoseChannel *pchan;

    for (pchan = ob->pose->chanbase.first; pchan; pchan = pchan->next) {
      if ((pchan->bone) && (arm->layer & pchan->bone->layer) && (pchan->mpath)) {
        /* new target for bone */
        mpt = MEM_callocN(sizeof(MPathTarget), "MPathTarget PoseBone");
        BLI_addtail(targets, mpt);

        mpt->mpath = pchan->mpath;
        mpt->ob = ob;
        mpt->pchan = pchan;
      }
    }
  }
}

/* ........ */

/* perform baking for the targets on the current frame */
static void motionpaths_calc_bake_targets(ListBase *targets, int cframe)
{
  MPathTarget *mpt;

  /* for each target, check if it can be baked on the current frame */
  for (mpt = targets->first; mpt; mpt = mpt->next) {
    bMotionPath *mpath = mpt->mpath;

    /* current frame must be within the range the cache works for
     * - is inclusive of the first frame, but not the last otherwise we get buffer overruns
     */
    if ((cframe < mpath->start_frame) || (cframe >= mpath->end_frame)) {
      continue;
    }

    /* get the relevant cache vert to write to */
    bMotionPathVert *mpv = mpath->points + (cframe - mpath->start_frame);

    Object *ob_eval = mpt->ob_eval;

    /* Lookup evaluated pose channel, here because the depsgraph
     * evaluation can change them so they are not cached in mpt. */
    bPoseChannel *pchan_eval = NULL;
    if (mpt->pchan) {
      pchan_eval = BKE_pose_channel_find_name(ob_eval->pose, mpt->pchan->name);
    }

    /* pose-channel or object path baking? */
    if (pchan_eval) {
      /* heads or tails */
      if (mpath->flag & MOTIONPATH_FLAG_BHEAD) {
        copy_v3_v3(mpv->co, pchan_eval->pose_head);
      }
      else {
        copy_v3_v3(mpv->co, pchan_eval->pose_tail);
      }

      /* result must be in worldspace */
      mul_m4_v3(ob_eval->obmat, mpv->co);
    }
    else {
      /* worldspace object location */
      copy_v3_v3(mpv->co, ob_eval->obmat[3]);
    }

    float mframe = (float)(cframe);

    /* Tag if it's a keyframe */
    if (BLI_dlrbTree_search_exact(&mpt->keys, compare_ak_cfraPtr, &mframe)) {
      mpv->flag |= MOTIONPATH_VERT_KEY;
    }
    else {
      mpv->flag &= ~MOTIONPATH_VERT_KEY;
    }

    /* Incremental update on evaluated object if possible, for fast updating
     * while dragging in transform. */
    bMotionPath *mpath_eval = NULL;
    if (mpt->pchan) {
      mpath_eval = (pchan_eval) ? pchan_eval->mpath : NULL;
    }
    else {
      mpath_eval = ob_eval->mpath;
    }

    if (mpath_eval && mpath_eval->length == mpath->length) {
      bMotionPathVert *mpv_eval = mpath_eval->points + (cframe - mpath_eval->start_frame);
      *mpv_eval = *mpv;

      GPU_VERTBUF_DISCARD_SAFE(mpath_eval->points_vbo);
      GPU_BATCH_DISCARD_SAFE(mpath_eval->batch_line);
      GPU_BATCH_DISCARD_SAFE(mpath_eval->batch_points);
    }
  }
}

/* Get pointer to animviz settings for the given target. */
static bAnimVizSettings *animviz_target_settings_get(MPathTarget *mpt)
{
  if (mpt->pchan != NULL) {
    return &mpt->ob->pose->avs;
  }
  return &mpt->ob->avs;
}

static void motionpath_get_global_framerange(ListBase *targets, int *r_sfra, int *r_efra)
{
  *r_sfra = INT_MAX;
  *r_efra = INT_MIN;
  LISTBASE_FOREACH (MPathTarget *, mpt, targets) {
    *r_sfra = min_ii(*r_sfra, mpt->mpath->start_frame);
    *r_efra = max_ii(*r_efra, mpt->mpath->end_frame);
  }
}

static int motionpath_get_prev_keyframe(MPathTarget *mpt, DLRBT_Tree *fcu_keys, int current_frame)
{
  if (current_frame <= mpt->mpath->start_frame) {
    return mpt->mpath->start_frame;
  }

  float current_frame_float = current_frame;
  DLRBT_Node *node = BLI_dlrbTree_search_prev(fcu_keys, compare_ak_cfraPtr, &current_frame_float);
  if (node == NULL) {
    return mpt->mpath->start_frame;
  }

  ActKeyColumn *key_data = (ActKeyColumn *)node;
  return key_data->cfra;
}

static int motionpath_get_prev_prev_keyframe(MPathTarget *mpt,
                                             DLRBT_Tree *fcu_keys,
                                             int current_frame)
{
  int frame = motionpath_get_prev_keyframe(mpt, fcu_keys, current_frame);
  return motionpath_get_prev_keyframe(mpt, fcu_keys, frame);
}

static int motionpath_get_next_keyframe(MPathTarget *mpt, DLRBT_Tree *fcu_keys, int current_frame)
{
  if (current_frame >= mpt->mpath->end_frame) {
    return mpt->mpath->end_frame;
  }

  float current_frame_float = current_frame;
  DLRBT_Node *node = BLI_dlrbTree_search_next(fcu_keys, compare_ak_cfraPtr, &current_frame_float);
  if (node == NULL) {
    return mpt->mpath->end_frame;
  }

  ActKeyColumn *key_data = (ActKeyColumn *)node;
  return key_data->cfra;
}

static int motionpath_get_next_next_keyframe(MPathTarget *mpt,
                                             DLRBT_Tree *fcu_keys,
                                             int current_frame)
{
  int frame = motionpath_get_next_keyframe(mpt, fcu_keys, current_frame);
  return motionpath_get_next_keyframe(mpt, fcu_keys, frame);
}

static bool motionpath_check_can_use_keyframe_range(MPathTarget *UNUSED(mpt),
                                                    AnimData *adt,
                                                    ListBase *fcurve_list)
{
  if (adt == NULL || fcurve_list == NULL) {
    return false;
  }
  /* NOTE: We might needed to do a full frame range update if there is a specific setup of NLA
   * or drivers or modifiers on the f-curves. */
  return true;
}

static void motionpath_calculate_update_range(MPathTarget *mpt,
                                              AnimData *adt,
                                              ListBase *fcurve_list,
                                              int current_frame,
                                              int *r_sfra,
                                              int *r_efra)
{
  *r_sfra = INT_MAX;
  *r_efra = INT_MIN;

  /* If the current frame is outside of the configured motion path range we ignore update of this
   * motion path by using invalid frame range where start frame is above the end frame. */
  if (current_frame < mpt->mpath->start_frame || current_frame > mpt->mpath->end_frame) {
    return;
  }

  /* Similar to the case when there is only a single keyframe: need to update en entire range to
   * a constant value. */
  if (!motionpath_check_can_use_keyframe_range(mpt, adt, fcurve_list)) {
    *r_sfra = mpt->mpath->start_frame;
    *r_efra = mpt->mpath->end_frame;
    return;
  }

  /* NOTE: Iterate over individual f-curves, and check their keyframes individually and pick a
   * widest range from them. This is because it's possible to have more narrow keyframe on a
   * channel which wasn't edited.
   * Could be optimized further by storing some flags about which channels has been modified so
   * we ignore all others (which can potentially make an update range unnecessary wide). */
  for (FCurve *fcu = fcurve_list->first; fcu != NULL; fcu = fcu->next) {
    DLRBT_Tree fcu_keys;
    BLI_dlrbTree_init(&fcu_keys);
    fcurve_to_keylist(adt, fcu, &fcu_keys, 0);

    int fcu_sfra = motionpath_get_prev_prev_keyframe(mpt, &fcu_keys, current_frame);
    int fcu_efra = motionpath_get_next_next_keyframe(mpt, &fcu_keys, current_frame);

    /* Extend range further, since acceleration compensation propagates even further away. */
    if (fcu->auto_smoothing != FCURVE_SMOOTH_NONE) {
      fcu_sfra = motionpath_get_prev_prev_keyframe(mpt, &fcu_keys, fcu_sfra);
      fcu_efra = motionpath_get_next_next_keyframe(mpt, &fcu_keys, fcu_efra);
    }

    if (fcu_sfra <= fcu_efra) {
      *r_sfra = min_ii(*r_sfra, fcu_sfra);
      *r_efra = max_ii(*r_efra, fcu_efra);
    }

    BLI_dlrbTree_free(&fcu_keys);
  }
}

static void motionpath_free_free_tree_data(ListBase *targets)
{
  LISTBASE_FOREACH (MPathTarget *, mpt, targets) {
    BLI_dlrbTree_free(&mpt->keys);
  }
}

/* Perform baking of the given object's and/or its bones' transforms to motion paths
 * - scene: current scene
 * - ob: object whose flagged motion-paths should get calculated
 * - recalc: whether we need to
 */
/* TODO: include reports pointer? */
void animviz_calc_motionpaths(Depsgraph *depsgraph,
                              Main *bmain,
                              Scene *scene,
                              ListBase *targets,
                              eAnimvizCalcRange range,
                              bool restore)
{
  /* Sanity check. */
  if (ELEM(NULL, targets, targets->first)) {
    return;
  }

  const int cfra = CFRA;
  int sfra = INT_MAX, efra = INT_MIN;
  switch (range) {
    case ANIMVIZ_CALC_RANGE_CURRENT_FRAME:
      motionpath_get_global_framerange(targets, &sfra, &efra);
      if (sfra > efra) {
        return;
      }
      if (cfra < sfra || cfra > efra) {
        return;
      }
      sfra = efra = cfra;
      break;
    case ANIMVIZ_CALC_RANGE_CHANGED:
      /* Nothing to do here, will be handled later when iterating through the targets. */
      break;
    case ANIMVIZ_CALC_RANGE_FULL:
      motionpath_get_global_framerange(targets, &sfra, &efra);
      if (sfra > efra) {
        return;
      }
      break;
  }

  /* get copies of objects/bones to get the calculated results from
   * (for copy-on-write evaluation), so that we actually get some results
   */

  /* TODO: Create a copy of background depsgraph that only contain these entities,
   * and only evaluates them.
   *
   * For until that is done we force dependency graph to not be active, so we don't lose unkeyed
   * changes during updating the motion path.
   * This still doesn't include unkeyed changes to the path itself, but allows to have updates in
   * an environment when auto-keying and pose paste is used. */

  const bool is_active_depsgraph = DEG_is_active(depsgraph);
  if (is_active_depsgraph) {
    DEG_make_inactive(depsgraph);
  }

  LISTBASE_FOREACH (MPathTarget *, mpt, targets) {
    mpt->ob_eval = DEG_get_evaluated_object(depsgraph, mpt->ob);

    AnimData *adt = BKE_animdata_from_id(&mpt->ob_eval->id);

    /* build list of all keyframes in active action for object or pchan */
    BLI_dlrbTree_init(&mpt->keys);

    ListBase *fcurve_list = NULL;
    if (adt) {
      /* get pointer to animviz settings for each target */
      bAnimVizSettings *avs = animviz_target_settings_get(mpt);

      /* it is assumed that keyframes for bones are all grouped in a single group
       * unless an option is set to always use the whole action
       */
      if ((mpt->pchan) && (avs->path_viewflag & MOTIONPATH_VIEW_KFACT) == 0) {
        bActionGroup *agrp = BKE_action_group_find_name(adt->action, mpt->pchan->name);

        if (agrp) {
          fcurve_list = &agrp->channels;
          agroup_to_keylist(adt, agrp, &mpt->keys, 0);
        }
      }
      else {
        fcurve_list = &adt->action->curves;
        action_to_keylist(adt, adt->action, &mpt->keys, 0);
      }
    }

    if (range == ANIMVIZ_CALC_RANGE_CHANGED) {
      int mpt_sfra, mpt_efra;
      motionpath_calculate_update_range(mpt, adt, fcurve_list, cfra, &mpt_sfra, &mpt_efra);
      if (mpt_sfra <= mpt_efra) {
        sfra = min_ii(sfra, mpt_sfra);
        efra = max_ii(efra, mpt_efra);
      }
    }
  }

  if (sfra > efra) {
    motionpath_free_free_tree_data(targets);
    return;
  }

  /* calculate path over requested range */
  CLOG_INFO(&LOG,
            1,
            "Calculating MotionPaths between frames %d - %d (%d frames)",
            sfra,
            efra,
            efra - sfra + 1);
  for (CFRA = sfra; CFRA <= efra; CFRA++) {
    if (range == ANIMVIZ_CALC_RANGE_CURRENT_FRAME) {
      /* For current frame, only update tagged. */
      BKE_scene_graph_update_tagged(depsgraph, bmain);
    }
    else {
      /* Update relevant data for new frame. */
      motionpaths_calc_update_scene(depsgraph);
    }

    /* perform baking for targets */
    motionpaths_calc_bake_targets(targets, CFRA);
  }

  /* reset original environment */
  /* NOTE: We don't always need to reevaluate the main scene, as the depsgraph
   * may be a temporary one that works on a subset of the data.
   * We always have to restore the current frame though. */
  CFRA = cfra;
  if (range != ANIMVIZ_CALC_RANGE_CURRENT_FRAME && restore) {
    motionpaths_calc_update_scene(depsgraph);
  }

  if (is_active_depsgraph) {
    DEG_make_active(depsgraph);
  }

  /* clear recalc flags from targets */
  LISTBASE_FOREACH (MPathTarget *, mpt, targets) {
    bMotionPath *mpath = mpt->mpath;

    /* get pointer to animviz settings for each target */
    bAnimVizSettings *avs = animviz_target_settings_get(mpt);

    /* clear the flag requesting recalculation of targets */
    avs->recalc &= ~ANIMVIZ_RECALC_PATHS;

    /* Clean temp data */
    BLI_dlrbTree_free(&mpt->keys);

    /* Free previous batches to force update. */
    GPU_VERTBUF_DISCARD_SAFE(mpath->points_vbo);
    GPU_BATCH_DISCARD_SAFE(mpath->batch_line);
    GPU_BATCH_DISCARD_SAFE(mpath->batch_points);
  }
}