# include "ohana.h"
# include <gsl/gsl_sf_legendre.h>

SHterms *SHtermsInit (int lmax) {

  // 0,0  -- N = 1, S = 1
  // 1,-1; 1,0; 1,1 -- N = 3, S = 4 (l+1)^2
  // 2,-2; 2,-1; 2,0; 2,1; 2,2 -- N = 5, S = 9 (l+1)^2
  // 3,-3; 3,-2; 3,-1; 3,0; 3,1; 3,2, 3,3 -- N = 7, S = 16 (l+1)^2

  SHterms *terms = NULL;
  ALLOCATE (terms, SHterms, 1);
  terms->lmax = lmax;
  terms->Nterms = SQ(lmax+1);

  ALLOCATE_ZERO (terms->Fr, double, terms->Nterms);
  ALLOCATE_ZERO (terms->Fi, double, terms->Nterms);
  ALLOCATE_ZERO (terms->l,     int, terms->Nterms);
  ALLOCATE_ZERO (terms->m,     int, terms->Nterms);
  
  int l, m;
  int n = 0;
  terms->l[n] = 0;
  terms->m[n] = 0;
  n++;
  for (l = 1; l <= terms->lmax; l++) {
    for (m = -l; m <= l; m++) {
      terms->l[n] = l;
      terms->m[n] = m;
      n++;
    }
  }
  return terms;
}

void SHtermsFree (SHterms *terms) {

  free (terms->Fr);
  free (terms->Fi);
  free (terms->l);
  free (terms->m);

  free (terms);
  return;
}

void SHtermsForRD (SHterms *terms, double R, double D) {

  // given a point on the sky R,D

  // for the real terms, Y_lm =
  // (m >  0) : Y_{l,m} = K(l,m) P_{l,m}*cos(m*R)    = gsl_sphPlm(l,|m|,cos(90-D))*cos(|m|*R)  [cos(x) == cos(-x)]
  // (m <  0) : Y_{l,m} = K(l,m) P_{l,m}*sin(|m|*R)  = gsl_sphPlm(l,|m|,cos(90-D))*sin(|m|*R)
  // (m == 0) : Y_{l,m} = K(l,m) P_{l,m}             = gsl_sphPlm(l,|m|,cos(90-D))

  double Rrad = R * RAD_DEG;
  double Drad = D * RAD_DEG;

  int l, m;

  int n = 0;
  for (l = 0; l <= terms->lmax; l++) {

    for (m = -l; m <= l; m++) {
    
      int m_abs = abs(m);
      double aP_lm = gsl_sf_legendre_sphPlm(l, m_abs, cos(M_PI/2 - Drad)); // alt-P_lm = P_lm (l, |m|, cos(90 - D))

      terms->Fr[n] = aP_lm * cos (m_abs * Rrad) ; 
      terms->Fi[n] = aP_lm * sin (m_abs * Rrad);
      n++;
    }
  }
}
/** I can confirm that these terms match those generated by V. Makarov's program **/

void SHtermsForLM (SHterms *terms, double R, double D, int l, int m) {

  double Rrad = R * RAD_DEG;
  double Drad = D * RAD_DEG;

  int m_abs = abs(m);
  double aP_lm = gsl_sf_legendre_sphPlm(l, m_abs, cos(M_PI/2 - Drad)); // alt-P_lm = P_lm (l, |m|, cos(90 - D))

  terms->Fr[0] = aP_lm * cos (m_abs * Rrad) ; 
  terms->Fi[0] = aP_lm * sin (m_abs * Rrad);
  terms->l[0] = l;
  terms->m[0] = m;
}

VSHterms *VSHtermsInit (int lmax) {

  // 0,0  -- N = 1, S = 1
  // 1,-1; 1,0; 1,1 -- N = 3, S = 4 (l+1)^2
  // 2,-2; 2,-1; 2,0; 2,1; 2,2 -- N = 5, S = 9 (l+1)^2
  // 3,-3; 3,-2; 3,-1; 3,0; 3,1; 3,2, 3,3 -- N = 7, S = 16 (l+1)^2

  VSHterms *terms = NULL;
  ALLOCATE (terms, VSHterms, 1);
  terms->lmax = lmax;
  terms->Nterms = SQ(lmax+1);

  ALLOCATE (terms->dR_B, double, terms->Nterms);
  ALLOCATE (terms->dR_E, double, terms->Nterms);
  ALLOCATE (terms->dD_B, double, terms->Nterms);
  ALLOCATE (terms->dD_E, double, terms->Nterms);
  ALLOCATE (terms->l,    int,    terms->Nterms);
  ALLOCATE (terms->m,    int,    terms->Nterms);
  
  int l, m;
  int n = 0;
  terms->l[n] = 0;
  terms->m[n] = 0;
  n++;
  for (l = 1; l <= terms->lmax; l++) {
    for (m = -l; m <= l; m++) {
      terms->l[n] = l;
      terms->m[n] = m;
      n++;
    }
  }

  return terms;
}

void VSHtermsFree (VSHterms *terms) {

  if (!terms) return;

  free (terms->dR_B);
  free (terms->dR_E);
  free (terms->dD_B);
  free (terms->dD_E);
  free (terms->l);
  free (terms->m);
  
  free (terms);

  return;
}

void VSHtermsForRD (VSHterms *terms, double R, double D) {

  // given a point on the sky R,D

  // we are calculating 2 terms: (dR,dD) = (dR^B + dR^E, dD^B + dD^E)

  // ^E terms are the 'electric' component, ^B are the 'magnetic' component

  // dR^B = \sum_{l,m} b_{l,m} dR^B_{l,m} (l,m,R,D) 
  // dD^B = \sum_{l,m} b_{l,m} dD^B_{l,m} (l,m,R,D)

  // dR^E = \sum_{l,m} e_{l,m} dR^E_{l,m} (l,m,R,D) 
  // dD^E = \sum_{l,m} e_{l,m} dD^E_{l,m} (l,m,R,D)
  
  // dR^E_{l,m} is Valeri's Elm, item[0], dD^E_{l,m} is Valeri's Elm, item[1]
  // dR^B_{l,m} is Valeri's Blm, item[0], dD^B_{l,m} is Valeri's Blm, item[1]

  // NOTE : as implemented by Valeri, the coeffs b_{l,m} are the same for both R and D
  // components; ditto for e_{l,m}

  // there are thus 4 X_l,m terms that need to be generated, for dR^B, dR^E, dD^B, dD^E.
  // Valeri calls these S_lm_lat and S_lm_lon

  // dR^B_{l,m} = +S^D_{l,m}(l,m,R,D)
  // dD^B_{l,m} = -S^R_{l,m}(l,m,R,D) / cos(D)

  // dR^E_{l,m} = +S^R_{l,m}(l,m,R,D) / cos(D)
  // dD^E_{l,m} = +S^D_{l,m}(l,m,R,D)

  // S^R_{l,m} == Slm_lng
  // (m >  0) : -m S_{l,m}(l,-m,R,D) 
  // (m <= 0) : -m S_{l,m}(l,-m,R,D)
  // NOTE: these are in fact the same...
  // NOTE: sending in -m means we get sin() for (m > 0) and cos() for (m <= 0)

  // S_{l,m}(l,m,R,D) :
  // aP_{l,m} = gsl_Plm(l,|m|,cos(90-D))

  // (m >  0) : S_{l,m} = aP_{l,m}*cos(m*R)    = gsl_Plm(l,|m|,cos(90-D))*cos(|m|*R)  [cos(x) == cos(-x)]
  // (m <  0) : S_{l,m} = aP_{l,m}*sin(|m|*R)  = gsl_Plm(l,|m|,cos(90-D))*sin(|m|*R)
  // (m == 0) : S_{l,m} = aP_{l,m}             = gsl_Plm(l,|m|,cos(90-D))

  // S^R_{l,m} = -m S_{l,m}(l,-m,R,D)
  // (m >  0) : S^R_{l,m} = -m aP_lm*sin(|m|*R)
  // (m <  0) : S^R_{l,m} = -m aP_lm*cos(|m|*R)
  // (m == 0) : S^R_{l,m} = 0

  // S_^D_{l,m} == Slm_lat:
  //
  // aP_{l,  m}(D) = gsl_Plm(l,  |m|,cos(90-D))
  // aP_{l+1,m}(D) = gsl_Plm(l+1,|m|,cos(90-D))

  // dPlm_dD = ((l+1)*sin(D)*aP_{l,m} - (l+1-|m|)*aP_{l+1,m}) / cos(D)
  // dSlm_dD
  // (m == 0) : dPlm_dD
  // (m > 0) : dPlm_dD*cos( m *R)
  // (m < 0) : dPlm_dD*sin(|m|*R)

  double Rrad = R * RAD_DEG;
  double Drad = D * RAD_DEG;

  int l, m;

  // Vlm(l=0) = 0, so assign directly (we ignore for fitting)
  int n = 0;
  terms->dR_B[n] = 0.0;
  terms->dD_B[n] = 0.0;
  terms->dR_E[n] = 0.0;
  terms->dD_E[n] = 0.0;
  n++;

  for (l = 1; l <= terms->lmax; l++) {

    for (m = -l; m <= l; m++) {
    
      int m_abs = abs(m);
      double aP_lm = gsl_sf_legendre_sphPlm(l, m_abs, cos(M_PI/2 - Drad)); // alt-P_lm = P_lm (l, |m|, cos(90 - D))

      // note that for m == 0, S_R_lm -> 0 (due to m*)
      double S_R_lm = (m > 0) ? -m * aP_lm * sin (m_abs * Rrad) : -m * aP_lm * cos (m_abs * Rrad) ; 
      
      // double aP_lp1m = gsl_sf_legendre_Plm(l + 1, m_abs, cos(M_PI/2 - Drad));
      double aP_lp1m = gsl_sf_legendre_sphPlm(l + 1, m_abs, cos(M_PI/2 - Drad));
      // reuse for aP_lm above?
      
      double dPlm_dD = ((l+1) * sin(Drad) * aP_lm - (l+1-m_abs)*aP_lp1m) / cos(Drad);
      // XXX : this diverges @ D = 90.0??
      // this is derived from the recursion relation:
      // (1 - x^2) (d/dx) P_l^m(x) = (l + 1) x P_l^m(x) - (l - m + 1) P^m_{l+1}
      // x = cos \theta = sin D
      // (1 - x^2) = sin^2 \theta = cos^2 D
      // (d/d\theta) P_l^m(x) = [(d/dx) P_l^m(x)] [dx/d\theta]
      // dx/d\theta = d cos\theta / d\theta = -sin \theta = -cos D [d/dD changes the sign]

      // (d/d\theta) P^m_l (x) = ((l + 1) sin D P_l^m(x) - (l - m + 1) P^m_{l+1}) / cos D

      double S_D_lm = (m < 0) ? dPlm_dD * sin(m_abs * Rrad) : dPlm_dD * cos(m_abs * Rrad);
  
      terms->dR_B[n] = +S_D_lm;
      terms->dD_B[n] = -S_R_lm / cos(Drad);
      terms->dR_E[n] = +S_R_lm / cos(Drad);
      terms->dD_E[n] = +S_D_lm;

      // fprintf (stderr, "%d : %d,%d : %f %f : %f %f %f %f\n", n, l, m, S_R_lm, S_D_lm, terms->dR_B[n], terms->dR_E[n], terms->dD_B[n], terms->dD_E[n]);
      // fprintf (stderr, "%d : %d,%d : %f %f %f %f\n", n-1, l, m, terms->dR_B[n], terms->dR_E[n], terms->dD_B[n], terms->dD_E[n]);
      n++;
    }
  }
}
/** I can confirm that these terms match those generated by V. Makarov's program **/

void VSHtermsForLM (VSHterms *terms, double R, double D, int l, int m) {

  double Rrad = R * RAD_DEG;
  double Drad = D * RAD_DEG;

  int m_abs = abs(m);
  double aP_lm = gsl_sf_legendre_Plm(l, m_abs, cos(M_PI/2 - Drad)); // alt-P_lm = P_lm (l, |m|, cos(90 - D))

  // note that for m == 0, S_R_lm -> 0 (due to m*)
  double S_R_lm = (m > 0) ? -m * aP_lm * sin (m_abs * Rrad) : -m * aP_lm * cos (m_abs * Rrad) ; 
      
  double aP_lp1m = gsl_sf_legendre_Plm(l + 1, m_abs, cos(M_PI/2 - Drad));
  // reuse for aP_lm above?
      
  double dPlm_dD = ((l+1) * sin(Drad) * aP_lm - (l+1-m_abs)*aP_lp1m) / cos(Drad);
  // XXX : this diverges @ D = 90.0??

  double S_D_lm = (m < 0) ? dPlm_dD * sin(m_abs * Rrad) : dPlm_dD * cos(m_abs * Rrad);
  
  terms->dR_B[0] = +S_D_lm;
  terms->dD_B[0] = -S_R_lm / cos(Drad);
  terms->dR_E[0] = +S_R_lm / cos(Drad);
  terms->dD_E[0] = +S_D_lm;
}
