/* Siconos is a program dedicated to modeling, simulation and control
 * of non smooth dynamical systems.
 *
 * Copyright 2020 INRIA.
 *
 * 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 <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <float.h>

#include "debug.h"
#include "NonSmoothDrivers.h"
#include "NCP_cst.h"
#include "NonlinearComplementarityProblem.h"
#include "NumericsMatrix.h"
#include "SolverOptions.h"

#define N 4

#ifdef __cplusplus
#undef restrict
#define restrict __restrict
#endif

typedef struct
{
  int id;
  double* xk;
  double h;
  double theta;
  double gamma;
  double g;
  double mu1;
  double kappa;
  unsigned int f_eval;
  unsigned int nabla_eval;
} data_ZIT;

static void compute_Fncp_ZIT(void* env, int n, double* restrict z, double* restrict F)
{
  data_ZIT* d = (data_ZIT*) env;
  double l0 = 2*z[0] - 1.0;
  double l1 = 2*z[2] - 1.0;
  double r1 = (d->mu1*l0 + l1)/(1.0 - d->kappa*l0*l1);
  double v_gamma = (d->xk[1] + d->gamma*(d->h*r1));
  double r0 = -d->kappa*l0*l1*(v_gamma);
  double v_theta = d->xk[1] + d->theta*(d->h*r1);
  F[0] = d->xk[0] + d->h*v_theta + d->h*r0 + z[1];
  F[2] = d->xk[1] + d->h*r1 + z[3];
  F[1] = 1.0 - z[0];
  F[3] = 1.0 - z[2];
  d->f_eval += 1;
}

static void compute_nabla_Fncp_ZIT(void* env, int n, double* restrict z, NumericsMatrix* restrict nabla_Fncp)
{
  double* restrict nabla_Fncp_dense = nabla_Fncp->matrix0;

  data_ZIT* d = (data_ZIT*) env;
  double l0 = 2.0*z[0] - 1.0;
  double l1 = 2.0*z[2] - 1.0;
  double invR = 1.0/(1.0 - d->kappa*l0*l1);
  double invR2 = invR*invR;
  double r1 = (d->mu1*l0 + l1)*invR;
  double v_gamma = d->xk[1] + d->gamma*(d->h*r1);

  nabla_Fncp_dense[2] = 2.0*d->h*(d->mu1+d->kappa*l1*l1)*invR2;
  nabla_Fncp_dense[2 + 4] = 0.0;
  nabla_Fncp_dense[2 + 2*4] = 2.0*d->h*(1.0 + d->kappa*l0*l0)*invR2;
  nabla_Fncp_dense[2 + 3*4] = 1.0;

  nabla_Fncp_dense[0] = -2.0*d->h*d->kappa*l1*(v_gamma) + d->h*(d->theta - d->gamma*d->kappa*l0*l1)*nabla_Fncp_dense[2];
  nabla_Fncp_dense[0 + 1*4] = 1.0;
  nabla_Fncp_dense[0 + 2*4] = -2.0*d->h*d->kappa*l0*(v_gamma) + d->h*(d->theta - d->gamma*d->kappa*l0*l1)*nabla_Fncp_dense[2 + 2*4];
  nabla_Fncp_dense[0 + 3*4] = 0.0;

  nabla_Fncp_dense[1] = -1.0;
  nabla_Fncp_dense[1 + 4] = 0.0;
  nabla_Fncp_dense[1 + 8] = 0.0;
  nabla_Fncp_dense[1 + 12] = 0.0;

  nabla_Fncp_dense[3] = 0.0;
  nabla_Fncp_dense[3 + 4] = 0.0;
  nabla_Fncp_dense[3 + 2*4] = -1.0;
  nabla_Fncp_dense[3 + 3*4] = 0.0;
  d->nabla_eval += 1;
}


int main(void)
{
  int info = 0;

  /* Set solver options */
  SolverOptions * options = solver_options_create(@PROBLEM_COLLECTION@);

  /* Create a NonlinearComplementarityProblem */
  NonlinearComplementarityProblem* problem = (NonlinearComplementarityProblem *)malloc(sizeof(NonlinearComplementarityProblem));

  problem->n = N;
  problem->compute_F = &compute_Fncp_ZIT;
  problem->compute_nabla_F = &compute_nabla_Fncp_ZIT;

  double xk[] = {1., 10.0};
  double t = 0.0;
  double T = 1.0;

  double z[N] = {0.};
  double F[N] = {0.};
  double nablaF[N*N] = {0.};

  data_ZIT sim_data;

  sim_data.id = -1;
  sim_data.xk = xk;
  sim_data.h = 1e-5;
  sim_data.theta = .5;
  sim_data.gamma = .5;
  sim_data.g = 9.81;
  sim_data.mu1 = 2.0;
  sim_data.kappa = .8;

 options->dparam[SICONOS_DPARAM_TOL] = N*DBL_EPSILON;
 options->iparam[SICONOS_IPARAM_PREALLOC] = 1;

  problem->env = &sim_data;
  problem->nabla_F = NM_create_from_data(NM_DENSE, N, N, nablaF);

  unsigned k = 0;

  while ((t <= T) && info == 0)
  {
    k++;
    info = ncp_driver(problem, z, F, &options);
    printf("step %d, iter = %d, error = %le\n", k, options->iparam[SICONOS_IPARAM_ITER_DONE], options->dparam[SICONOS_DPARAM_RESIDU]);
    if (info > 0)
    {
      printf("NCP_ZI1 info = %d\n", info);
      /* do some magic shit to force the algo to fix a solution */
      double x1 = F[0] - z[1];
      double x2 = F[2] - z[3];
      if ((x1 < -1e-7) && (fabs(z[1]) < 1e-10))
      {
        z[1] = -x1;
        z[0] = 1.0;
      }
      if ((xk[1] < -1e-7) && (fabs(z[3]) < 1e-10))
      {
        z[3] = -x2;
        z[2] = 1.0;
      }
      if (z[1] < -1e-7)
      {
        z[1] = 0.0;
        z[0] = 0.0;
      }
      if (z[3] < -1e-7)
      {
        z[3] = 0.0;
        z[2] = 0.0;
      }
      if ((z[1] > 1e-7) && (z[0] < 1.0 - 1e-7))
      {
        z[0] = 1.0;
      }
      if ((z[3] > 1e-7) && (z[2] < 1.0 - 1e-7))
      {
        z[2] = 1.0;
      }

      info = ncp_driver(problem, z, F, &options);

      if (info > 0)
      {
        printf("step %d, iter = %d, error = %le\n", k, options->iparam[SICONOS_IPARAM_ITER_DONE], options->dparam[SICONOS_DPARAM_RESIDU]);
        printf("NCP_ZI1 with solver @PROBLEM_COLLECTION@ failed completly\n");
        printf("NCP_ZIT1 failure at xk[0] =  %.*e; xk[1] =  %.*e\n", DECIMAL_DIG, xk[0], DECIMAL_DIG, xk[1]);
      }

    }
    xk[0] = F[0] - z[1];
    xk[1] = F[2] - z[3];
    t = k*sim_data.h;
  }
  solver_options_delete(options);
  free(problem->nabla_F);
  free(problem);
  return info;
}
