text: /*
  Copyright (c) by respective owners including Yahoo!, Microsoft, and
  individual contributors. All rights reserved.    Released under a BSD (revised)
  license as described in the file LICENSE.
  */
  #include <string>
  #include "gd.h"
  #include "vw.h"
  #include "rand48.h"
  #include "reductions.h"
  #include <math.h>
  
  using namespace std;
  using namespace LEARNER;
  
  #define NORM2 (m+1)
  
  struct update_data
  {
    struct OjaNewton *ON;
    float g;
    float sketch_cnt;
    float norm2_x;
    float *Zx;
    float *AZx;
    float *delta;
    float bdelta;
    float prediction;
  };
  
  struct OjaNewton
  {
    vw* all;
    int m;
    int epoch_size;
    float alpha;
    int cnt;
    int t;
  
    float *ev;
    float *b;
    float *D;
    float **A;
    float **K;
  
    float *zv;
    float *vv;
    float *tmp;
  
    example **buffer;
    float *weight_buffer;
    struct update_data data;
  
    float learning_rate_cnt;
    bool normalize;
    bool random_init;
  
    void initialize_Z(parameters& weights)
    {
      uint32_t length = 1 << all->num_bits;
      if (normalize)   // initialize normalization part
      {
        for (uint32_t i = 0; i < length; i++)
          (&(weights.strided_index(i)))[NORM2] = 0.1f;
      }
      if(!random_init)
      {
        // simple initialization
        for (int i = 1; i <= m; i++)
          (&(weights.strided_index(i)))[i] = 1.f;
      }
      else
      {
        // more complicated initialization: orthgonal basis of a random matrix
  
        const double PI2 = 2.f * 3.1415927f;
  
        for (uint32_t i = 0; i < length; i++)
        {
          weight& w = weights.strided_index(i);
          for (int j = 1; j <= m; j++)
          {
            float r1 = merand48(all->random_state);
            float r2 = merand48(all->random_state);
            (&w)[j] = sqrt(-2.f * log(r1)) * (float)cos(PI2 * r2);
          }
        }
      }
  
      // Gram-Schmidt
      for (int j = 1; j <= m; j++)
      {
        for (int k = 1; k <= j - 1; k++)
        {
          double tmp = 0;
  
          for (uint32_t i = 0; i < length; i++)
            tmp += ((double)(&(weights.strided_index(i)))[j]) * (&(weights.strided_index(i)))[k];
          for (uint32_t i = 0; i < length; i++)
            (&(weights.strided_index(i)))[j] -= (float)tmp *(&(weights.strided_index(i)))[k];
        }
        double norm = 0;
        for (uint32_t i = 0; i < length; i++)
          norm += ((double)(&(weights.strided_index(i)))[j]) * (&(weights.strided_index(i)))[j];
        norm = sqrt(norm);
        for (uint32_t i = 0; i < length; i++)
          (&(weights.strided_index(i)))[j] /= (float)norm;
      }
    }
  
    void compute_AZx()
    {
      for (int i = 1; i <= m; i++)
      {
        data.AZx[i] = 0;
        for (int j = 1; j <= i; j++)
        {
          data.AZx[i] += A[i][j] * data.Zx[j];
        }
      }
    }
  
    void update_eigenvalues()
    {
      for (int i = 1; i <= m; i++)
      {
        float gamma = fmin(learning_rate_cnt / t, 1.f);
        float tmp = data.AZx[i] * data.sketch_cnt;
  
        if (t == 1)
        {
          ev[i] = gamma * tmp * tmp;
        }
        else
        {
          ev[i] = (1 - gamma) * t * ev[i] / (t - 1) + gamma * t * tmp * tmp;
        }
      }
    }
  
    void compute_delta()
    {
      data.bdelta = 0;
      for (int i = 1; i <= m; i++)
      {
        float gamma = fmin(learning_rate_cnt / t, 1.f);
  
        // if different learning rates are used
        /*data.delta[i] = gamma * data.AZx[i] * data.sketch_cnt;
        for (int j = 1; j < i; j++) {
            data.delta[i] -= A[i][j] * data.delta[j];
        }
        data.delta[i] /= A[i][i];*/
  
        // if a same learning rate is used
        data.delta[i] = gamma * data.Zx[i] * data.sketch_cnt;
  
        data.bdelta += data.delta[i] * b[i];
      }
    }
  
    void update_K()
    {
      float tmp = data.norm2_x * data.sketch_cnt * data.sketch_cnt;
      for (int i = 1; i <= m; i++)
      {
        for (int j = 1; j <= m; j++)
        {
          K[i][j] += data.delta[i] * data.Zx[j] * data.sketch_cnt;
          K[i][j] += data.delta[j] * data.Zx[i] * data.sketch_cnt;
          K[i][j] += data.delta[i] * data.delta[j] * tmp;
        }
      }
    }
  
    void update_A()
    {
      for (int i = 1; i <= m; i++)
      {
  
        for (int j = 1; j < i; j++)
        {
          zv[j] = 0;
          for (int k = 1; k <= i; k++)
          {
            zv[j] += A[i][k] * K[k][j];
          }
        }
  
        for (int j = 1; j < i; j++)
        {
          vv[j] = 0;
          for (int k = 1; k <= j; k++)
          {
            vv[j] += A[j][k] * zv[k];
          }
        }
  
        for (int j = 1; j < i; j++)
        {
          for (int k = j; k < i; k++)
          {
            A[i][j] -= vv[k] * A[k][j];
          }
        }
  
        float norm = 0;
        for (int j = 1; j <= i; j++)
        {
          float temp = 0;
          for (int k = 1; k <= i; k++)
          {
            temp += K[j][k] * A[i][k];
          }
          norm += A[i][j]*temp;
        }
        norm = sqrtf(norm);
  
        for (int j = 1; j <= i; j++)
        {
          A[i][j] /= norm;
        }
      }
    }
  
    void update_b()
    {
      for (int j = 1; j <= m; j++)
      {
        float tmp = 0;
        for (int i = j; i <= m; i++)
        {
          tmp += ev[i] * data.AZx[i] * A[i][j] / (alpha * (alpha + ev[i]));
        }
        b[j] += tmp * data.g;
      }
    }
  
    void update_D()
    {
      for (int j = 1; j <= m; j++)
      {
        float scale = fabs(A[j][j]);
        for (int i = j+1; i <= m; i++)
          scale = fmin(fabs(A[i][j]), scale);
        if (scale < 1e-10) continue;
        for (int i = 1; i <= m; i++)
        {
          A[i][j] /= scale;
          K[j][i] *= scale;
          K[i][j] *= scale;
        }
        b[j] /= scale;
        D[j] *= scale;
        //printf("D[%d] = %f\n", j, D[j]);
      }
    }
  
    void check()
    {
      double max_norm = 0;
      for (int i = 1; i <= m; i++)
        for (int j = i; j <= m ; j++)
          max_norm = fmax(max_norm, fabs(K[i][j]));
      //printf("|K| = %f\n", max_norm);
      if (max_norm < 1e7) return;
  
      // implicit -> explicit representation
      // printf("begin conversion: t = %d, norm(K) = %f\n", t, max_norm);
  
      // first step: K <- AKA'
  
      // K <- AK
      for (int j = 1; j <= m; j++)
      {
        memset(tmp, 0, sizeof(double) * (m+1));
  
        for (int i = 1; i <= m; i++)
        {
          for (int h = 1; h <= m; h++)
          {
            tmp[i] += A[i][h] * K[h][j];
          }
        }
  
        for (int i = 1; i <= m; i++)
          K[i][j] = tmp[i];
      }
      // K <- KA'
      for (int i = 1; i <= m; i++)
      {
        memset(tmp, 0, sizeof(double) * (m+1));
  
        for (int j = 1; j <= m; j++)
          for (int h = 1; h <= m; h++)
            tmp[j] += K[i][h] * A[j][h];
  
        for (int j = 1; j <= m; j++)
        {
          K[i][j] = tmp[j];
        }
      }
  
      //second step: w[0] <- w[0] + (DZ)'b, b <- 0.
  
      uint32_t length = 1 << all->num_bits;
      for (uint32_t i = 0; i < length; i++)
      {
        weight& w = all->weights.strided_index(i);
        for (int j = 1; j <= m; j++)
          w += (&w)[j] * b[j] * D[j];
      }
  
      memset(b, 0, sizeof(double) * (m+1));
  
      //third step: Z <- ADZ, A, D <- Identity
  
      //double norm = 0;
      for (uint32_t i = 0; i < length; ++i)
      {
        memset(tmp, 0, sizeof(float) * (m+1));
        weight& w = all->weights.strided_index(i);
        for (int j = 1; j <= m; j++)
        {
          for (int h = 1; h <= m; ++h)
            tmp[j] += A[j][h] * D[h] * (&w)[h];
        }
        for (int j = 1; j <= m; ++j)
        {
          //norm = max(norm, fabs(tmp[j]));
          (&w)[j] = tmp[j];
        }
      }
      //printf("|Z| = %f\n", norm);
  
      for (int i = 1; i <= m; i++)
      {
        memset(A[i], 0, sizeof(double) * (m+1));
        D[i] = 1;
        A[i][i] = 1;
      }
    }
  };
  
  void keep_example(vw& all, OjaNewton& ON, example& ec)
  {
    output_and_account_example(all, ec);
  }
  
  void finish(OjaNewton& ON)
  {
    free(ON.ev);
    free(ON.b);
    free(ON.D);
    free(ON.buffer);
    free(ON.weight_buffer);
    free(ON.zv);
    free(ON.vv);
    free(ON.tmp);
  
    for (int i = 1; i <= ON.m; i++)
    {
      free(ON.A[i]);
      free(ON.K[i]);
    }
    free(ON.A);
    free(ON.K);
  
    free(ON.data.Zx);
    free(ON.data.AZx);
    free(ON.data.delta);
  }
  
  void make_pred(update_data& data, float x, float& wref)
  {
    int m = data.ON->m;
    float* w = &wref;
  
    if (data.ON->normalize)
    {
      x /= sqrt(w[NORM2]);
    }
  
    data.prediction += w[0] * x;
    for (int i = 1; i <= m; i++)
    {
      data.prediction += w[i] * x * data.ON->D[i] * data.ON->b[i];
    }
  }
  
  void predict(OjaNewton& ON, base_learner&, example& ec)
  {
    ON.data.prediction = 0;
    GD::foreach_feature<update_data, make_pred>(*ON.all, ec, ON.data);
    ec.partial_prediction = (float)ON.data.prediction;
    ec.pred.scalar = GD::finalize_prediction(ON.all->sd, ec.partial_prediction);
  }
  
  void update_Z_and_wbar(update_data& data, float x, float& wref)
  {
    float* w = &wref;
    int m = data.ON->m;
    if (data.ON->normalize) x /= sqrt(w[NORM2]);
    float s = data.sketch_cnt * x;
  
    for (int i = 1; i <= m; i++)
    {
      w[i] += data.delta[i] * s / data.ON->D[i];
    }
    w[0] -= s * data.bdelta;
  }
  
  void compute_Zx_and_norm(update_data& data, float x, float& wref)
  {
    float* w = &wref;
    int m = data.ON->m;
    if (data.ON->normalize) x /= sqrt(w[NORM2]);
  
    for (int i = 1; i <= m; i++)
    {
      data.Zx[i] += w[i] * x * data.ON->D[i];
    }
    data.norm2_x += x * x;
  }
  
  void update_wbar_and_Zx(update_data& data, float x, float& wref)
  {
    float* w = &wref;
    int m = data.ON->m;
    if (data.ON->normalize) x /= sqrt(w[NORM2]);
  
    float g = data.g * x;
  
    for (int i = 1; i <= m; i++)
    {
      data.Zx[i] += w[i] * x * data.ON->D[i];
    }
    w[0] -= g / data.ON->alpha;
  }
  
  
  void update_normalization(update_data& data, float x, float& wref)
  {
    float* w = &wref;
    int m = data.ON->m;
  
    w[NORM2] += x * x * data.g * data.g;
  }
  
  void learn(OjaNewton& ON, base_learner& base, example& ec)
  {
    assert(ec.in_use);
  
    // predict
    predict(ON, base, ec);
  
    update_data& data = ON.data;
    data.g = ON.all->loss->first_derivative(ON.all->sd, ec.pred.scalar, ec.l.simple.label)*ec.l.simple.weight;
    data.g /= 2; // for half square loss
  
    if(ON.normalize) GD::foreach_feature<update_data, update_normalization>(*ON.all, ec, data);
  
    ON.buffer[ON.cnt] = &ec;
    ON.weight_buffer[ON.cnt++] = data.g / 2;
  
    if (ON.cnt == ON.epoch_size)
    {
      for (int k = 0; k < ON.epoch_size; k++, ON.t++)
      {
        example& ex = *(ON.buffer[k]);
        data.sketch_cnt = ON.weight_buffer[k];
  
        data.norm2_x = 0;
        memset(data.Zx, 0, sizeof(float)* (ON.m+1));
        GD::foreach_feature<update_data, compute_Zx_and_norm>(*ON.all, ex, data);
        ON.compute_AZx();
  
        ON.update_eigenvalues();
        ON.compute_delta();
  
        ON.update_K();
  
        GD::foreach_feature<update_data, update_Z_and_wbar>(*ON.all, ex, data);
      }
  
      ON.update_A();
      //ON.update_D();
    }
  
    memset(data.Zx, 0, sizeof(float)* (ON.m+1));
    GD::foreach_feature<update_data, update_wbar_and_Zx>(*ON.all, ec, data);
    ON.compute_AZx();
  
    ON.update_b();
    ON.check();
  
    if (ON.cnt == ON.epoch_size)
    {
      ON.cnt = 0;
      for (int k = 0; k < ON.epoch_size; k++)
      {
        VW::finish_example(*ON.all, *ON.buffer[k]);
      }
    }
  }
  
  
  void save_load(OjaNewton& ON, io_buf& model_file, bool read, bool text)
  {
    vw& all = *ON.all;
    if (read)
    {
      initialize_regressor(all);
      ON.initialize_Z(all.weights);
    }
  
    if (model_file.files.size() > 0)
    {
      bool resume = all.save_resume;
      stringstream msg;
      msg << ":"<< resume <<"\n";
      bin_text_read_write_fixed(model_file, (char *)&resume, sizeof (resume), "", read, msg, text);
  
  
      if (resume)
        GD::save_load_online_state(all, model_file, read, text);
      else
        GD::save_load_regressor(all, model_file, read, text);
    }
  }
  
  base_learner* OjaNewton_setup(arguments& arg)
  {
    auto ON = scoped_calloc_or_throw<OjaNewton>();
    if (arg.new_options("OjaNewton options")
        .critical("OjaNewton", "Online Newton with Oja's Sketch")
        ("sketch_size", ON->m, 10, "size of sketch")
        ("epoch_size", ON->epoch_size, 1, "size of epoch")
        ("alpha", ON->alpha, 1.f, "mutiplicative constant for indentiy")
        ("alpha_inverse", po::value<float>(), "one over alpha, similar to learning rate")
        ("learning_rate_cnt", ON->learning_rate_cnt, 2.f, "constant for the learning rate 1/t")
        ("normalize", ON->normalize, true, "normalize the features or not")
        ("random_init", ON->random_init, true, "randomize initialization of Oja or not").missing())
      return nullptr;
  
    ON->all = arg.all;
  
    if (arg.vm.count("alpha_inverse"))
      ON->alpha = 1.f / arg.vm["alpha_inverse"].as<float>();
  
    ON->cnt = 0;
    ON->t = 1;
    ON->ev = calloc_or_throw<float>(ON->m+1);
    ON->b = calloc_or_throw<float>(ON->m+1);
    ON->D = calloc_or_throw<float>(ON->m+1);
    ON->A = calloc_or_throw<float*>(ON->m+1);
    ON->K = calloc_or_throw<float*>(ON->m+1);
    for (int i = 1; i <= ON->m; i++)
    {
      ON->A[i] = calloc_or_throw<float>(ON->m+1);
      ON->K[i] = calloc_or_throw<float>(ON->m+1);
      ON->A[i][i] = 1;
      ON->K[i][i] = 1;
      ON->D[i] = 1;
    }
  
    ON->buffer = calloc_or_throw<example*>(ON->epoch_size);
    ON->weight_buffer = calloc_or_throw<float>(ON->epoch_size);
  
    ON->zv = calloc_or_throw<float>(ON->m+1);
    ON->vv = calloc_or_throw<float>(ON->m+1);
    ON->tmp = calloc_or_throw<float>(ON->m+1);
  
    ON->data.ON = ON.get();
    ON->data.Zx = calloc_or_throw<float>(ON->m+1);
    ON->data.AZx = calloc_or_throw<float>(ON->m+1);
    ON->data.delta = calloc_or_throw<float>(ON->m+1);
  
    arg.all->weights.stride_shift((uint32_t)ceil(log2(ON->m + 2)));
  
    learner<OjaNewton, example>& l = init_learner(ON, learn, predict, arg.all->weights.stride());
    l.set_save_load(save_load);
    l.set_finish_example(keep_example);
    l.set_finish(finish);
    return make_base(l);
  }
  