#include "data.table.h"

SEXP fifelseR(SEXP l, SEXP a, SEXP b, SEXP na) {
  if (!isLogical(l)) {
    error(_("Argument 'test' must be logical."));
  }
  if (  (isS4(a) && !INHERITS(a, char_nanotime))
     || (isS4(b) && !INHERITS(b, char_nanotime)) ) {
    error(_("S4 class objects (except nanotime) are not supported."));
  }
  const int64_t len0 = xlength(l);
  const int64_t len1 = xlength(a);
  const int64_t len2 = xlength(b);
  const int64_t len3 = xlength(na);
  SEXPTYPE ta = TYPEOF(a);
  SEXPTYPE tb = TYPEOF(b);
  SEXPTYPE tn = TYPEOF(na);
  // na_a/b/n means a scalar NA (or NULL for the na argument), which is considered to be coerced into other types
  bool na_a = len1==1 && ta==LGLSXP && LOGICAL(a)[0]==NA_LOGICAL;
  bool na_b = len2==1 && tb==LGLSXP && LOGICAL(b)[0]==NA_LOGICAL;
  bool na_n = isNull(na) || (len3==1 && tn==LGLSXP && LOGICAL(na)[0]==NA_LOGICAL);

  if (!na_a && len1!=1 && len1!=len0)
    error(_("Length of 'yes' is %"PRId64" but must be 1 or length of 'test' (%"PRId64")."), len1, len0);
  if (!na_b && len2!=1 && len2!=len0)
    error(_("Length of 'no' is %"PRId64" but must be 1 or length of 'test' (%"PRId64")."), len2, len0);
  if (!na_n && len3!=1 && len3!=len0)
    error(_("Length of 'na' is %"PRId64" but must be 1 or length of 'test' (%"PRId64")."), len3, len0);
  
  int nprotect = 0;
  SEXPTYPE tans = !na_a ? ta : !na_b ? tb : !na_n ? tn : LGLSXP;
  if (!(na_a && na_b && na_n)) {
    SEXPTYPE ta0 = ta, tb0 = tb, tn0 = tn; // record the original type for error message use
    if (!na_b && tans==INTSXP && tb==REALSXP) tans = tb;
    if (!na_n && tans==INTSXP && tn==REALSXP) tans = tn;
    if (!na_a && tans==REALSXP && ta==INTSXP) {
      a = PROTECT(coerceVector(a, REALSXP)); nprotect++;
      ta = REALSXP;
    }
    // it's not possible that non-NA `yes`' type will be different from `tans`
    if (!na_b && tans==REALSXP && tb==INTSXP) {
      b = PROTECT(coerceVector(b, REALSXP)); nprotect++;
      tb = REALSXP;
    }
    if (!na_b && tans != tb) 
      error(_("'no' is of type %s but '%s' is %s. Please make all arguments have the same type."), type2char(tb0), tans==ta0 ? "yes" : "na", tans==ta0 ? type2char(ta0) : type2char(tn0));
    if (!na_n && tans==REALSXP && tn==INTSXP) {
      na = PROTECT(coerceVector(na, REALSXP)); nprotect++;
      tn = REALSXP;
    }
    if (!na_n && tans != tn) 
      error(_("'na' is of type %s but '%s' is %s. Please make all arguments have the same type."), type2char(tn0), tans==ta0 ? "yes" : "no", tans==ta0 ? type2char(ta0) : type2char(tb0));
  }
  
  if (!na_a && !na_b) {
    if (!R_compute_identical(PROTECT(getAttrib(a,R_ClassSymbol)), PROTECT(getAttrib(b,R_ClassSymbol)), 0))
      error(_("'yes' has different class than 'no'. Please make sure that both arguments have the same class."));
    UNPROTECT(2);
  }
  if (!na_a && !na_n) {
    if (!R_compute_identical(PROTECT(getAttrib(a,R_ClassSymbol)), PROTECT(getAttrib(na,R_ClassSymbol)), 0))
      error(_("'yes' has different class than 'na'. Please make sure that both arguments have the same class."));
    UNPROTECT(2);
  }
  if (!na_b && !na_n) {
    if (!R_compute_identical(PROTECT(getAttrib(b,R_ClassSymbol)), PROTECT(getAttrib(na,R_ClassSymbol)), 0))
      error(_("'no' has different class than 'na'. Please make sure that both arguments have the same class."));
    UNPROTECT(2);
  }
  
  if (isFactor(a) || isFactor(b)) {
    if (!na_a && !na_b) {
      if (!R_compute_identical(PROTECT(getAttrib(a,R_LevelsSymbol)), PROTECT(getAttrib(b,R_LevelsSymbol)), 0))
        error(_("'yes' and 'no' are both type factor but their levels are different."));
      UNPROTECT(2);  
    }
    if (!na_a && !na_n) {
      if (!R_compute_identical(PROTECT(getAttrib(a,R_LevelsSymbol)), PROTECT(getAttrib(na,R_LevelsSymbol)), 0))
        error(_("'yes' and 'na' are both type factor but their levels are different."));
      UNPROTECT(2);
    }
    if (!na_b && !na_n) {
      if (!R_compute_identical(PROTECT(getAttrib(b,R_LevelsSymbol)), PROTECT(getAttrib(na,R_LevelsSymbol)), 0))
        error(_("'no' and 'na' are both type factor but their levels are different."));
      UNPROTECT(2);
    }
  }

  const int64_t amask = len1>1 ? INT64_MAX : 0; // for scalar 'a' bitwise AND will reset iterator to first element: pa[i & amask] -> pa[0]
  const int64_t bmask = len2>1 ? INT64_MAX : 0;
  const int64_t nmask = len3>1 ? INT64_MAX : 0;

  const int *restrict pl = LOGICAL(l);
  SEXP ans = PROTECT(allocVector(tans, len0)); nprotect++;
  if (!na_a) 
    copyMostAttrib(a, ans);
  else if (!na_b)
    copyMostAttrib(b, ans);
  else if (!na_n)
    copyMostAttrib(na, ans);

  switch(tans) {
  case LGLSXP: {
    int *restrict pans = LOGICAL(ans);
    const int *restrict pa = na_a ? NULL : LOGICAL(a);
    const int *restrict pb = na_b ? NULL : LOGICAL(b);
    const int *restrict pna = na_n ? NULL : LOGICAL(na);
    const int na = NA_LOGICAL;
    #pragma omp parallel for num_threads(getDTthreads(len0, true))
    for (int64_t i=0; i<len0; ++i) {
      pans[i] = pl[i]==0 ?
                (na_b ? na : pb[i & bmask]) :
                pl[i]==1 ? 
                (na_a ? na : pa[i & amask]) : 
                (na_n ? na : pna[i & nmask]);
    }
  } break;
  case INTSXP: {
    int *restrict pans = INTEGER(ans);
    const int *restrict pa = na_a ? NULL : INTEGER(a);
    const int *restrict pb = na_b ? NULL : INTEGER(b);
    const int *restrict pna = na_n ? NULL : INTEGER(na);
    const int na = NA_INTEGER;
    #pragma omp parallel for num_threads(getDTthreads(len0, true))
    for (int64_t i=0; i<len0; ++i) {
      pans[i] = pl[i]==0 ?
                (na_b ? na : pb[i & bmask]) :
                pl[i]==1 ? 
                (na_a ? na : pa[i & amask]) : 
                (na_n ? na : pna[i & nmask]);
    }
  } break;
  case REALSXP: {
    double *restrict pans = REAL(ans);
    const double *restrict pa = na_a ? NULL : REAL(a);
    const double *restrict pb = na_b ? NULL : REAL(b);
    const double *restrict pna = na_n ? NULL : REAL(na);
    const double na = INHERITS(a, char_integer64) ? NA_INT64_D : NA_REAL;
    #pragma omp parallel for num_threads(getDTthreads(len0, true))
    for (int64_t i=0; i<len0; ++i) {
      pans[i] = pl[i]==0 ?
                (na_b ? na : pb[i & bmask]) :
                pl[i]==1 ? 
                (na_a ? na : pa[i & amask]) : 
                (na_n ? na : pna[i & nmask]);
    }
  } break;
  case STRSXP : {
    const SEXP *restrict pa = na_a ? NULL : STRING_PTR_RO(a);
    const SEXP *restrict pb = na_b ? NULL : STRING_PTR_RO(b);
    const SEXP *restrict pna = na_n ? NULL : STRING_PTR_RO(na);
    const SEXP na = NA_STRING;
    for (int64_t i=0; i<len0; ++i) {
      SET_STRING_ELT(
        ans, i, pl[i]==0 ?
                (na_b ? na : pb[i & bmask]) :
                pl[i]==1 ? 
                (na_a ? na : pa[i & amask]) : 
                (na_n ? na : pna[i & nmask])
      );
    }
  } break;
  case CPLXSXP : {
    Rcomplex *restrict pans = COMPLEX(ans);
    const Rcomplex *restrict pa = na_a ? NULL : COMPLEX(a);
    const Rcomplex *restrict pb = na_b ? NULL : COMPLEX(b);
    const Rcomplex *restrict pna = na_n ? NULL : COMPLEX(na);
    const Rcomplex na = NA_CPLX;
    #pragma omp parallel for num_threads(getDTthreads(len0, true))
    for (int64_t i=0; i<len0; ++i) {
      pans[i] = pl[i]==0 ?
                (na_b ? na : pb[i & bmask]) :
                pl[i]==1 ? 
                (na_a ? na : pa[i & amask]) : 
                (na_n ? na : pna[i & nmask]);
    }
  } break;
  case VECSXP : {
    const SEXP *restrict pa = na_a ? NULL : SEXPPTR_RO(a);
    const SEXP *restrict pb = na_b ? NULL : SEXPPTR_RO(b);
    const SEXP *restrict pna = na_n ? NULL : SEXPPTR_RO(na);
    for (int64_t i=0; i<len0; ++i) {
      if (pl[i] == NA_LOGICAL) {
        if (!na_n) SET_VECTOR_ELT(ans, i, pna[i & nmask]);
      } else if (pl[i]==0) {
        if (!na_b) SET_VECTOR_ELT(ans, i, pb[i & bmask]);
      } else if (pl[i]==1) {
        if (!na_a) SET_VECTOR_ELT(ans, i, pa[i & amask]);
      }
    }
  } break;
  default:
    error(_("Type '%s' is not supported"), type2char(ta));
  }

  SEXP l_names = PROTECT(getAttrib(l, R_NamesSymbol)); nprotect++;
  if (!isNull(l_names))
    setAttrib(ans, R_NamesSymbol, l_names);

  UNPROTECT(nprotect);
  return ans;
}

SEXP fcaseR(SEXP rho, SEXP args) {
  const int narg=length(args); // `default` will take the last two positions
  if (narg % 2) {
    error(_("Received %d inputs; please supply an even number of arguments in ..., "
            "consisting of logical condition, resulting value pairs (in that order). "
            "Note that the default argument must be named explicitly, e.g., default=0"), narg - 2);
  }
  int nprotect=0, l;
  int64_t len0=0, len1=0, len2=0;
  SEXP ans=R_NilValue, value0=R_NilValue, tracker=R_NilValue, whens=R_NilValue, thens=R_NilValue;
  PROTECT_INDEX Iwhens, Ithens;
  PROTECT_WITH_INDEX(whens, &Iwhens); nprotect++;
  PROTECT_WITH_INDEX(thens, &Ithens); nprotect++;
  SEXPTYPE type0=NILSXP;
  // naout means if the output is scalar logic na
  bool imask = true, naout = false, idefault = false;
  int *restrict p = NULL;
  const int n = narg/2;
  for (int i=0; i<n; ++i) {
    idefault = i == (n - 1); // mark if the current eval is the `default` on R side
    REPROTECT(whens = eval(SEXPPTR_RO(args)[2*i], rho), Iwhens);
    REPROTECT(thens = eval(SEXPPTR_RO(args)[2*i+1], rho), Ithens);
    if (isS4(thens) && !INHERITS(thens, char_nanotime)) {
      error(_("S4 class objects (except nanotime) are not supported. Please see https://github.com/Rdatatable/data.table/issues/4131."));
    }
    if (!isLogical(whens)) {
      error(_("Argument #%d must be logical but was of type %s."), 2*i+1, type2char(TYPEOF(whens)));
    }
    const int *restrict pwhens = LOGICAL(whens);
    l = 0;
    if (i == 0) {
      len0 = xlength(whens);
      len2 = len0;
      type0 = TYPEOF(thens);
      value0 = thens;
      ans = PROTECT(allocVector(type0, len0)); nprotect++;
      copyMostAttrib(thens, ans);
      tracker = PROTECT(allocVector(INTSXP, len0)); nprotect++;
      p = INTEGER(tracker);     
    } else {
      imask = false;
      naout = xlength(thens) == 1 && TYPEOF(thens) == LGLSXP && LOGICAL(thens)[0]==NA_LOGICAL;
      if (xlength(whens) != len0) {
        // no need to check `idefault` here because the con for default is always `TRUE`
        error(_("Argument #%d has length %lld which differs from that of argument #1 (%lld). "
                "Please make sure all logical conditions have the same length."),
                i*2+1, (long long)xlength(whens), (long long)len0);
      }
      if (!naout && TYPEOF(thens) != type0) {
        if (idefault) {
          error(_("Resulting value is of type %s but 'default' is of type %s. "
                  "Please make sure that both arguments have the same type."), type2char(type0), type2char(TYPEOF(thens)));
        } else {
          error(_("Argument #%d is of type %s, however argument #2 is of type %s. "
                  "Please make sure all output values have the same type."),
                  i*2+2, type2char(TYPEOF(thens)), type2char(type0));
        }
      }
      if (!naout) {
        if (!R_compute_identical(PROTECT(getAttrib(value0, R_ClassSymbol)),  PROTECT(getAttrib(thens, R_ClassSymbol)), 0)) {
          if (idefault) {
            error(_("Resulting value has different class than 'default'. "
                    "Please make sure that both arguments have the same class."));
          } else {
            error(_("Argument #%d has different class than argument #2, "
                    "Please make sure all output values have the same class."), i*2+2);
          }
        }
        UNPROTECT(2); // class(value0), class(thens)
      }
      if (!naout && isFactor(value0)) {
        if (!R_compute_identical(PROTECT(getAttrib(value0, R_LevelsSymbol)),  PROTECT(getAttrib(thens, R_LevelsSymbol)), 0)) {
          if (idefault) {
            error(_("Resulting value and 'default' are both type factor but their levels are different."));
          } else {
            error(_("Argument #2 and argument #%d are both factor but their levels are different."), i*2+2);
          }
        }
        UNPROTECT(2); // levels(value0), levels(thens)
      }
    }
    len1 = xlength(thens);
    if (len1 != len0 && len1 != 1) {
      if (idefault) {
        error(_("Length of 'default' must be 1 or %lld."), (long long)len0);
      } else {
        error(_("Length of output value #%d (%lld) must either be 1 or match the length of the logical condition (%lld)."), i*2+2, (long long)len1, (long long)len0);
      }
    }
    int64_t thenMask = len1>1 ? INT64_MAX : 0;
    switch(TYPEOF(ans)) {
    case LGLSXP: {
      const int *restrict pthens;
      if (!naout) pthens = LOGICAL(thens); // the content is not useful if out is NA_LOGICAL scalar
      int *restrict pans = LOGICAL(ans);
      const int pna = NA_LOGICAL;
      for (int64_t j=0; j<len2; ++j) {
        const int64_t idx = imask ? j : p[j];
        if (pwhens[idx]==1) {
          pans[idx] = naout ? pna : pthens[idx & thenMask];
        } else {
          if (imask) {
            pans[j] = pna;
          }
          p[l++] = idx;
        }
      }
    } break;
    case INTSXP: {
      const int *restrict pthens;
      if (!naout) pthens = INTEGER(thens); // the content is not useful if out is NA_LOGICAL scalar
      int *restrict pans = INTEGER(ans);
      const int pna = NA_INTEGER;
      for (int64_t j=0; j<len2; ++j) {
        const int64_t idx = imask ? j : p[j];
        if (pwhens[idx]==1) {
          pans[idx] = naout ? pna : pthens[idx & thenMask];
        } else {
          if (imask) {
            pans[j] = pna;
          }
          p[l++] = idx;
        }
      }
    } break;
    case REALSXP: {
      const double *restrict pthens;
      if (!naout) pthens = REAL(thens); // the content is not useful if out is NA_LOGICAL scalar
      double *restrict pans = REAL(ans);
      const double na_double = INHERITS(ans, char_integer64) ? NA_INT64_D : NA_REAL;
      const double pna = na_double;
      for (int64_t j=0; j<len2; ++j) {
        const int64_t idx = imask ? j : p[j];
        if (pwhens[idx]==1) {
          pans[idx] = naout ? pna : pthens[idx & thenMask];
        } else {
          if (imask) {
            pans[j] = pna;
          }
          p[l++] = idx;
        }
      }
    } break;
    case CPLXSXP: {
      const Rcomplex *restrict pthens;
      if (!naout) pthens = COMPLEX(thens); // the content is not useful if out is NA_LOGICAL scalar
      Rcomplex *restrict pans = COMPLEX(ans);
      const Rcomplex pna = NA_CPLX;
      for (int64_t j=0; j<len2; ++j) {
        const int64_t idx = imask ? j : p[j];
        if (pwhens[idx]==1) {
          pans[idx] = naout ? pna : pthens[idx & thenMask];
        } else {
          if (imask) {
            pans[j] = pna;
          }
          p[l++] = idx;
        }
      }
    } break;
    case STRSXP: {
      const SEXP *restrict pthens;
      if (!naout) pthens = STRING_PTR_RO(thens); // the content is not useful if out is NA_LOGICAL scalar
      const SEXP pna = NA_STRING;
      for (int64_t j=0; j<len2; ++j) {
        const int64_t idx = imask ? j : p[j];
        if (pwhens[idx]==1) {
          SET_STRING_ELT(ans, idx, naout ? pna : pthens[idx & thenMask]);
        } else {
          if (imask) {
            SET_STRING_ELT(ans, j, pna);
          }
          p[l++] = idx;
        }
      }
    } break;
    case VECSXP: {
      // the default value of VECSXP is `NULL` so we don't need to explicitly
      // assign the NA values as it does for other atomic types
      const SEXP *restrict pthens;
      if (!naout) pthens = SEXPPTR_RO(thens); // the content is not useful if out is NA_LOGICAL scalar
      for (int64_t j=0; j<len2; ++j) {
        const int64_t idx = imask ? j : p[j];
        if (pwhens[idx]==1) {
          if (!naout) SET_VECTOR_ELT(ans, idx, pthens[idx & thenMask]);
        } else {
          p[l++] = idx;
        }
      }
    } break;
    default:
      error(_("Type '%s' is not supported."), type2char(TYPEOF(ans)));
    }
    if (l==0) {
      break;  // stop early as nothing left to do
    }
    len2 = l;
  }
  UNPROTECT(nprotect); // whens, thens, ans, tracker
  return ans;
}
