#pragma once
#include <stdio.h>
#include "../vowpalwabbit/parser.h"
#include "../vowpalwabbit/vw.h"

typedef uint32_t fid;

struct vw_namespace
{ char namespace_letter;
public: vw_namespace(const char c) : namespace_letter(c) {}
};

class ezexample
{
private:
  vw*vw_ref;
  vw*vw_par_ref;   // an extra parser if we're multithreaded
  bool is_multiline;

  char str[2];
  example*ec;
  bool we_create_ec;
  std::vector<fid> past_seeds;
  fid current_seed;
  size_t quadratic_features_num;
  float quadratic_features_sqr;
  char current_ns;
  bool ns_exists[256];
  bool example_changed_since_prediction;

  v_array<example*> example_copies;

  ezexample(const ezexample & ex);
  ezexample & operator=(const ezexample & ex);

  example* get_new_example()
  { example* new_ec = VW::new_unused_example(*vw_par_ref);
    vw_par_ref->p->lp.default_label(&new_ec->l);
    new_ec->tag.erase();
    new_ec->indices.erase();
    for (size_t i=0; i<256; i++)
      new_ec->feature_space[i].erase();

    new_ec->ft_offset = 0;
    new_ec->num_features = 0;
    new_ec->partial_prediction = 0.;
    new_ec->updated_prediction = 0.;
    new_ec->passthrough = nullptr;
    new_ec->loss = 0.;
    new_ec->total_sum_feat_sq = 0.;
    new_ec->confidence = 0.;
    return new_ec;
  }

  void setup_new_ezexample(vw*this_vw, bool multiline, vw*this_vw_parser)
  { vw_ref = this_vw;
    vw_par_ref = (this_vw_parser == nullptr) ? this_vw : this_vw_parser;
    is_multiline = multiline;

    str[0] = 0; str[1] = 0;
    current_seed = 0;
    current_ns = 0;

    quadratic_features_num = 0;
    quadratic_features_sqr = 0.;

    for (size_t i=0; i<256; i++) ns_exists[i] = false;

    example_changed_since_prediction = true;
  }


  void setup_for_predict()
  { static example* empty_example = is_multiline ? VW::read_example(*vw_par_ref, (char*)"") : nullptr;
    if (example_changed_since_prediction)
    { mini_setup_example();
      vw_ref->learn(ec);
      if (is_multiline) vw_ref->learn(empty_example);
      example_changed_since_prediction = false;
    }
  }

public:

  // REAL FUNCTIONALITY
  // create a new ezexample by asking the vw parser for an example
  ezexample(vw*this_vw, bool multiline=false, vw*this_vw_parser=nullptr)
  { setup_new_ezexample(this_vw, multiline, this_vw_parser);
    example_copies = v_init<example*>();
    ec = get_new_example();
    we_create_ec = true;

    if (vw_ref->add_constant)
      VW::add_constant_feature(*vw_ref, ec);
  }

  // create a new ezexample by wrapping around an already existing example
  // we do NOT copy your data, therefore, WARNING:
  //   do NOT touch the underlying example unless you really know what you're done)
  ezexample(vw*this_vw, example*this_ec, bool multiline=false, vw*this_vw_parser=nullptr)
  { setup_new_ezexample(this_vw, multiline, this_vw_parser);

    ec = this_ec;
    we_create_ec = false;

    for (auto ns : ec->indices)
      ns_exists[ns] = true;
    if (current_ns != 0)
    { str[0] = current_ns;
      current_seed = VW::hash_space(*vw_ref, str);
    }
  }

  ~ezexample()   // calls finish_example *only* if we created our own example!
  { if (ec->in_use && VW::is_ring_example(*vw_par_ref, ec))
      VW::finish_example(*vw_par_ref, ec);
    for (auto ecc : example_copies)
      if (ecc->in_use && VW::is_ring_example(*vw_par_ref, ec))
        VW::finish_example(*vw_par_ref, ecc);
    example_copies.erase();
    free(example_copies.begin());
  }

  bool ensure_ns_exists(char c)    // returns TRUE iff we should ignore it :)
  { if (vw_ref->ignore_some && vw_ref->ignore[(int)c]) return true;
    if (ns_exists[(int)c]) return false;
    ec->indices.push_back((size_t)c);
    ns_exists[(int)c] = true;
    return false;
  }

  void addns(char c)
  { if (ensure_ns_exists(c)) return;

    ec->feature_space[(int)c].erase();
    past_seeds.push_back(current_seed);
    current_ns = c;
    str[0] = c;
    current_seed = VW::hash_space(*vw_ref, str);
  }

  void remns()
  { if (ec->indices.size() == 0)
    { current_seed = 0;
      current_ns = 0;
    }
    else
    { if (ns_exists[(int)current_ns])
      { ec->total_sum_feat_sq -= ec->feature_space[(int)current_ns].sum_feat_sq;
        ec->feature_space[(int)current_ns].erase();
        ec->num_features -= ec->feature_space[(int)current_ns].size();

        ns_exists[(int)current_ns] = false;
      }

      current_seed = past_seeds.back();
      past_seeds.pop_back();
      ec->indices.pop();
      example_changed_since_prediction = true;
    }
  }


  inline fid addf(char to_ns, fid fint, float v)
  { if (to_ns == 0) return 0;
    if (ensure_ns_exists(to_ns)) return 0;

    ec->feature_space[(int)to_ns].push_back(v, fint << vw_ref->weights.stride_shift());
    ec->total_sum_feat_sq += v * v;
    ec->num_features++;
    example_changed_since_prediction = true;
    return fint;
  }

  inline fid addf(fid fint, float v) { return addf(current_ns, fint, v); }

  // copy an entire namespace from this other example, you can even give it a new namespace name if you want!
  void add_other_example_ns(example& other, char other_ns, char to_ns)
  { if (ensure_ns_exists(to_ns)) return;
    features& fs = other.feature_space[(int)other_ns];
    for (size_t i = 0; i < fs.size(); i++)
      ec->feature_space[(int)to_ns].push_back(fs.values[i], fs.indicies[i]);
    ec->total_sum_feat_sq += fs.sum_feat_sq;
    ec->num_features += fs.size();
    example_changed_since_prediction = true;
  }
  void add_other_example_ns(example& other, char ns)    // default to_ns to other_ns
  { add_other_example_ns(other, ns, ns);
  }

  void add_other_example_ns(ezexample& other, char other_ns, char to_ns) { add_other_example_ns(*other.ec, other_ns, to_ns); }
  void add_other_example_ns(ezexample& other, char ns                  ) { add_other_example_ns(*other.ec, ns); }

  inline ezexample& set_label(std::string label)
  { VW::parse_example_label(*vw_par_ref, *ec, label);
    example_changed_since_prediction = true;
    return *this;
  }

  void mini_setup_example()
  { ec->partial_prediction = 0.;
    ec->weight = vw_par_ref->p->lp.get_weight(&ec->l);

    ec->num_features      -= quadratic_features_num;
    ec->total_sum_feat_sq -= quadratic_features_sqr;

    quadratic_features_num = 0;
    quadratic_features_sqr = 0.;

    for (std::vector<std::string>::iterator i = vw_ref->pairs.begin(); i != vw_ref->pairs.end(); i++)
    { quadratic_features_num
      += ec->feature_space[(int)(*i)[0]].size() * ec->feature_space[(int)(*i)[1]].size();
      quadratic_features_sqr
      += ec->feature_space[(int)(*i)[0]].sum_feat_sq * ec->feature_space[(int)(*i)[1]].sum_feat_sq;
    }
    ec->num_features      += quadratic_features_num;
    ec->total_sum_feat_sq += quadratic_features_sqr;
  }

  size_t get_num_features() { return ec->num_features; }

  example* get()
  { if (example_changed_since_prediction)
      mini_setup_example();
    return ec;
  }

  float predict()
  { setup_for_predict();
    return ec->pred.scalar;
  }

  float predict_partial()
  { setup_for_predict();
    return ec->partial_prediction;
  }

  void train()    // if multiline, add to stack; otherwise, actually train
  { if (example_changed_since_prediction)
    { mini_setup_example();
      example_changed_since_prediction = false;
    }

    if (!is_multiline)
    { vw_ref->learn(ec);
    }
    else       // is multiline
    { // we need to make a copy
      example* copy = get_new_example();
      assert(ec->in_use);
      VW::copy_example_data(vw_ref->audit, copy, ec, vw_par_ref->p->lp.label_size, vw_par_ref->p->lp.copy_label);
      assert(copy->in_use);
      vw_ref->learn(copy);
      example_copies.push_back(copy);
    }
  }

  void clear_features()
  { for (size_t i=0; i<256; i++)
    { if (current_ns == 0) break;
      remns();
    }
  }

  void finish()
  { static example* empty_example = is_multiline ? VW::read_example(*vw_par_ref, (char*)"") : nullptr;
    if (is_multiline)
    { vw_ref->learn(empty_example);
      for (auto ecc : example_copies)
        if (ecc->in_use)
          VW::finish_example(*vw_par_ref, ecc);
      example_copies.erase();
    }
  }


  // HELPER FUNCTIONALITY

  inline fid hash(std::string fstr)         { return VW::hash_feature(*vw_ref, fstr, current_seed); }
  inline fid hash(char*  fstr)         { return VW::hash_feature_cstr(*vw_ref, fstr, current_seed); }
  inline fid hash(char c, std::string fstr) { str[0] = c; return VW::hash_feature(*vw_ref, fstr, VW::hash_space(*vw_ref, str)); }
  inline fid hash(char c, char*  fstr) { str[0] = c; return VW::hash_feature_cstr(*vw_ref, fstr, VW::hash_space(*vw_ref, str)); }

  inline fid addf(fid    fint           ) { return addf(fint      , 1.0); }
  inline fid addf(std::string fstr, float val) { return addf(hash(fstr), val); }
  inline fid addf(std::string fstr           ) { return addf(hash(fstr), 1.0); }

  inline fid addf(char ns, fid    fint           ) { return addf(ns, fint          , 1.0); }
  inline fid addf(char ns, std::string fstr, float val) { return addf(ns, hash(ns, fstr), val); }
  inline fid addf(char ns, std::string fstr           ) { return addf(ns, hash(ns, fstr), 1.0); }

  inline ezexample& operator()(const vw_namespace&n) { addns(n.namespace_letter); return *this; }

  inline ezexample& operator()(fid         fint           ) { addf(fint, 1.0); return *this; }
  inline ezexample& operator()(std::string      fstr           ) { addf(fstr, 1.0); return *this; }
  inline ezexample& operator()(const char* fstr           ) { addf(fstr, 1.0); return *this; }
  inline ezexample& operator()(fid         fint, float val) { addf(fint, val); return *this; }
  inline ezexample& operator()(std::string      fstr, float val) { addf(fstr, val); return *this; }
  inline ezexample& operator()(const char* fstr, float val) { addf(fstr, val); return *this; }

  inline ezexample& operator()(char ns, fid         fint           ) { addf(ns, fint, 1.0); return *this; }
  inline ezexample& operator()(char ns, std::string      fstr           ) { addf(ns, fstr, 1.0); return *this; }
  inline ezexample& operator()(char ns, const char* fstr           ) { addf(ns, fstr, 1.0); return *this; }
  inline ezexample& operator()(char ns, fid         fint, float val) { addf(ns, fint, val); return *this; }
  inline ezexample& operator()(char ns, std::string      fstr, float val) { addf(ns, fstr, val); return *this; }
  inline ezexample& operator()(char ns, const char* fstr, float val) { addf(ns, fstr, val); return *this; }

  inline ezexample& operator()(  example&other, char other_ns, char to_ns) { add_other_example_ns(other, other_ns, to_ns); return *this; }
  inline ezexample& operator()(  example&other, char ns                  ) { add_other_example_ns(other, ns);              return *this; }
  inline ezexample& operator()(ezexample&other, char other_ns, char to_ns) { add_other_example_ns(other, other_ns, to_ns); return *this; }
  inline ezexample& operator()(ezexample&other, char ns                  ) { add_other_example_ns(other, ns);              return *this; }


  inline ezexample& operator--() { remns(); return *this; }

  inline float      operator()() { return predict(); }
};
