/* Copyright (C) 2015 Atsushi Togo */
/* All rights reserved. */

/* This file is part of phonopy. */

/* Redistribution and use in source and binary forms, with or without */
/* modification, are permitted provided that the following conditions */
/* are met: */

/* * Redistributions of source code must retain the above copyright */
/*   notice, this list of conditions and the following disclaimer. */

/* * Redistributions in binary form must reproduce the above copyright */
/*   notice, this list of conditions and the following disclaimer in */
/*   the documentation and/or other materials provided with the */
/*   distribution. */

/* * Neither the name of the phonopy project nor the names of its */
/*   contributors may be used to endorse or promote products derived */
/*   from this software without specific prior written permission. */

/* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS */
/* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT */
/* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS */
/* FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE */
/* COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, */
/* INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, */
/* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; */
/* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER */
/* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT */
/* LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN */
/* ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE */
/* POSSIBILITY OF SUCH DAMAGE. */

#include <stdio.h>
#include <stdlib.h>
#include <phonoc_array.h>
#include <phonoc_const.h>
#include <phonon3_h/interaction.h>
#include <phonon3_h/real_to_reciprocal.h>
#include <phonon3_h/reciprocal_to_normal.h>
#include <lapack_wrapper.h>

static const int index_exchange[6][3] = {{0, 1, 2},
                                         {2, 0, 1},
                                         {1, 2, 0},
                                         {2, 1, 0},
                                         {0, 2, 1},
                                         {1, 0, 2}};
static void real_to_normal(double *fc3_normal_squared,
                           PHPYCONST int (*g_pos)[4],
                           const int num_g_pos,
                           const double *freqs0,
                           const double *freqs1,
                           const double *freqs2,
                           const lapack_complex_double *eigvecs0,
                           const lapack_complex_double *eigvecs1,
                           const lapack_complex_double *eigvecs2,
                           const double *fc3,
                           const double q[9], /* q0, q1, q2 */
                           const Darray *shortest_vectors,
                           const int *multiplicity,
                           const double *masses,
                           const int *p2s_map,
                           const int *s2p_map,
                           const int *band_indices,
                           const int num_band0,
                           const int num_band,
                           const double cutoff_frequency,
                           const int triplet_index,
                           const int num_triplets,
                           const int openmp_at_bands);
static void real_to_normal_sym_q(double *fc3_normal_squared,
                                 PHPYCONST int (*g_pos)[4],
                                 const int num_g_pos,
                                 PHPYCONST double *freqs[3],
                                 PHPYCONST lapack_complex_double *eigvecs[3],
                                 const double *fc3,
                                 const double q[9], /* q0, q1, q2 */
                                 const Darray *shortest_vectors,
                                 const int *multiplicity,
                                 const double *masses,
                                 const int *p2s_map,
                                 const int *s2p_map,
                                 const int *band_indices,
                                 const int num_band0,
                                 const int num_band,
                                 const double cutoff_frequency,
                                 const int triplet_index,
                                 const int num_triplets,
                                 const int openmp_at_bands);

/* fc3_normal_squared[num_triplets, num_band0, num_band, num_band] */
void get_interaction(Darray *fc3_normal_squared,
                     const char *g_zero,
                     const Darray *frequencies,
                     const Carray *eigenvectors,
                     const Iarray *triplets,
                     const int *grid_address,
                     const int *mesh,
                     const Darray *fc3,
                     const Darray *shortest_vectors,
                     const Iarray *multiplicity,
                     const double *masses,
                     const int *p2s_map,
                     const int *s2p_map,
                     const int *band_indices,
                     const int symmetrize_fc3_q,
                     const double cutoff_frequency)
{
  int i, num_band, num_band0, num_band_prod, openmp_per_triplets;
  int j, k, l, jkl, num_g_pos;
  int (*g_pos)[4];

  g_pos = NULL;

  num_band0 = fc3_normal_squared->dims[1];
  num_band = frequencies->dims[1];
  num_band_prod = num_band0 * num_band * num_band;

  if (triplets->dims[0] > num_band) {
    openmp_per_triplets = 1;
  } else {
    openmp_per_triplets = 0;
  }

#pragma omp parallel for schedule(guided) private(j, k, l, jkl, num_g_pos, g_pos) if (openmp_per_triplets)
  for (i = 0; i < triplets->dims[0]; i++) {
    num_g_pos = 0;
    jkl = 0;
    g_pos = (int(*)[4])malloc(sizeof(int[4]) * num_band_prod);
    for (j = 0; j < num_band0; j++) {
      for (k = 0; k < num_band; k++) {
        for (l = 0; l < num_band; l++) {
          if (!g_zero[jkl + i * num_band_prod]) {
            g_pos[num_g_pos][0] = j;
            g_pos[num_g_pos][1] = k;
            g_pos[num_g_pos][2] = l;
            g_pos[num_g_pos][3] = jkl;
            num_g_pos++;
          }
          fc3_normal_squared->data[jkl + i * num_band_prod] = 0;
          jkl++;
        }
      }
    }

    get_interaction_at_triplet(
      fc3_normal_squared->data + i * num_band_prod,
      num_band0,
      num_band,
      g_pos,
      num_g_pos,
      frequencies->data,
      eigenvectors->data,
      triplets->data + i * 3,
      grid_address,
      mesh,
      fc3->data,
      shortest_vectors,
      multiplicity->data,
      masses,
      p2s_map,
      s2p_map,
      band_indices,
      symmetrize_fc3_q,
      cutoff_frequency,
      i,
      triplets->dims[0],
      1 - openmp_per_triplets);

    free(g_pos);
    g_pos = NULL;
  }
}

void get_interaction_at_triplet(double *fc3_normal_squared,
                                const int num_band0,
                                const int num_band,
                                PHPYCONST int (*g_pos)[4],
                                const int num_g_pos,
                                const double *frequencies,
                                const lapack_complex_double *eigenvectors,
                                const int *triplet,
                                const int *grid_address,
                                const int *mesh,
                                const double *fc3,
                                const Darray *shortest_vectors,
                                const int *multiplicity,
                                const double *masses,
                                const int *p2s_map,
                                const int *s2p_map,
                                const int *band_indices,
                                const int symmetrize_fc3_q,
                                const double cutoff_frequency,
                                const int triplet_index, /* only for print */
                                const int num_triplets, /* only for print */
                                const int openmp_at_bands)
{
  int j, k;
  double *freqs[3];
  lapack_complex_double *eigvecs[3];
  double q[9];

  for (j = 0; j < 3; j++) {
    for (k = 0; k < 3; k++) {
      q[j * 3 + k] = ((double)grid_address[triplet[j] * 3 + k]) / mesh[k];
    }
  }

  if (symmetrize_fc3_q) {
    for (j = 0; j < 3; j++) {
      freqs[j] = (double*)malloc(sizeof(double) * num_band);
      eigvecs[j] = (lapack_complex_double*)
        malloc(sizeof(lapack_complex_double) * num_band * num_band);
      for (k = 0; k < num_band; k++) {
        freqs[j][k] = frequencies[triplet[j] * num_band + k];
      }
      for (k = 0; k < num_band * num_band; k++) {
        eigvecs[j][k] = eigenvectors[triplet[j] * num_band * num_band + k];
      }
    }
    real_to_normal_sym_q(fc3_normal_squared,
                         g_pos,
                         num_g_pos,
                         freqs,
                         eigvecs,
                         fc3,
                         q, /* q0, q1, q2 */
                         shortest_vectors,
                         multiplicity,
                         masses,
                         p2s_map,
                         s2p_map,
                         band_indices,
                         num_band0,
                         num_band,
                         cutoff_frequency,
                         triplet_index,
                         num_triplets,
                         openmp_at_bands);
    for (j = 0; j < 3; j++) {
      free(freqs[j]);
      freqs[j] = NULL;
      free(eigvecs[j]);
      eigvecs[j] = NULL;
    }
  } else {
    real_to_normal(fc3_normal_squared,
                   g_pos,
                   num_g_pos,
                   frequencies + triplet[0] * num_band,
                   frequencies + triplet[1] * num_band,
                   frequencies + triplet[2] * num_band,
                   eigenvectors + triplet[0] * num_band * num_band,
                   eigenvectors + triplet[1] * num_band * num_band,
                   eigenvectors + triplet[2] * num_band * num_band,
                   fc3,
                   q, /* q0, q1, q2 */
                   shortest_vectors,
                   multiplicity,
                   masses,
                   p2s_map,
                   s2p_map,
                   band_indices,
                   num_band0,
                   num_band,
                   cutoff_frequency,
                   triplet_index,
                   num_triplets,
                   openmp_at_bands);
  }
}

static void real_to_normal(double *fc3_normal_squared,
                           PHPYCONST int (*g_pos)[4],
                           const int num_g_pos,
                           const double *freqs0,
                           const double *freqs1,
                           const double *freqs2,
                           const lapack_complex_double *eigvecs0,
                           const lapack_complex_double *eigvecs1,
                           const lapack_complex_double *eigvecs2,
                           const double *fc3,
                           const double q[9], /* q0, q1, q2 */
                           const Darray *shortest_vectors,
                           const int *multiplicity,
                           const double *masses,
                           const int *p2s_map,
                           const int *s2p_map,
                           const int *band_indices,
                           const int num_band0,
                           const int num_band,
                           const double cutoff_frequency,
                           const int triplet_index,
                           const int num_triplets,
                           const int openmp_at_bands)
{
  int num_patom;
  lapack_complex_double *fc3_reciprocal;

  num_patom = num_band / 3;

  fc3_reciprocal =
    (lapack_complex_double*)malloc(sizeof(lapack_complex_double) *
                                   num_patom * num_patom * num_patom * 27);
  real_to_reciprocal(fc3_reciprocal,
                     q,
                     fc3,
                     shortest_vectors,
                     multiplicity,
                     p2s_map,
                     s2p_map,
                     openmp_at_bands);

#ifdef MEASURE_R2N
  if (openmp_at_bands && num_triplets > 0) {
    printf("At triplet %d/%d (# of bands=%d):\n",
           triplet_index, num_triplets, num_band0);
  }
#endif
  reciprocal_to_normal_squared(fc3_normal_squared,
                               g_pos,
                               num_g_pos,
                               fc3_reciprocal,
                               freqs0,
                               freqs1,
                               freqs2,
                               eigvecs0,
                               eigvecs1,
                               eigvecs2,
                               masses,
                               band_indices,
                               num_band0,
                               num_band,
                               cutoff_frequency,
                               openmp_at_bands);

  free(fc3_reciprocal);
  fc3_reciprocal = NULL;
}

static void real_to_normal_sym_q(double *fc3_normal_squared,
                                 PHPYCONST int (*g_pos)[4],
                                 const int num_g_pos,
                                 PHPYCONST double *freqs[3],
                                 PHPYCONST lapack_complex_double *eigvecs[3],
                                 const double *fc3,
                                 const double q[9], /* q0, q1, q2 */
                                 const Darray *shortest_vectors,
                                 const int *multiplicity,
                                 const double *masses,
                                 const int *p2s_map,
                                 const int *s2p_map,
                                 const int *band_indices,
                                 const int num_band0,
                                 const int num_band,
                                 const double cutoff_frequency,
                                 const int triplet_index,
                                 const int num_triplets,
                                 const int openmp_at_bands)
{
  int i, j, k, l;
  int band_ex[3];
  double q_ex[9];
  double *fc3_normal_squared_ex;

  fc3_normal_squared_ex =
    (double*)malloc(sizeof(double) * num_band * num_band * num_band);

  for (i = 0; i < num_band0 * num_band * num_band; i++) {
    fc3_normal_squared[i] = 0;
  }

  for (i = 0; i < 6; i++) {
    for (j = 0; j < 3; j ++) {
      for (k = 0; k < 3; k ++) {
        q_ex[j * 3 + k] = q[index_exchange[i][j] * 3 + k];
      }
    }
    real_to_normal(fc3_normal_squared_ex,
                   g_pos,
                   num_g_pos,
                   freqs[index_exchange[i][0]],
                   freqs[index_exchange[i][1]],
                   freqs[index_exchange[i][2]],
                   eigvecs[index_exchange[i][0]],
                   eigvecs[index_exchange[i][1]],
                   eigvecs[index_exchange[i][2]],
                   fc3,
                   q_ex, /* q0, q1, q2 */
                   shortest_vectors,
                   multiplicity,
                   masses,
                   p2s_map,
                   s2p_map,
                   band_indices,
                   num_band,
                   num_band,
                   cutoff_frequency,
                   triplet_index,
                   num_triplets,
                   openmp_at_bands);
    for (j = 0; j < num_band0; j++) {
      for (k = 0; k < num_band; k++) {
        for (l = 0; l < num_band; l++) {
          band_ex[0] = band_indices[j];
          band_ex[1] = k;
          band_ex[2] = l;
          fc3_normal_squared[j * num_band * num_band +
                             k * num_band +
                             l] +=
            fc3_normal_squared_ex[band_ex[index_exchange[i][0]] *
                                  num_band * num_band +
                                  band_ex[index_exchange[i][1]] * num_band +
                                  band_ex[index_exchange[i][2]]] / 6;
        }
      }
    }
  }

  free(fc3_normal_squared_ex);

}
