import _ from 'lodash';
import { lookupGrowthStandards, round } from './utils';

type zscoreFun = (y: number, m: number, l: number, s:number) => number

function calc_sd(sd: number, m: number, l: number, s: number): number {
  return m * ((1 + l * s * sd) ** (1 / l));
}
// ' Implements the non adjusted z-score
// '
// ' @param y y
// ' @param m m
// ' @param l l
// ' @param s s
// '
// ' @references
// ' http://www.who.int/childgrowth/standards/Chap_7.pdf?ua=1
// '
// ' @noRd
function compute_zscore(y: number, m: number, l: number, s: number): number {
  return ((y / m) ** l - 1) / (s * l);
}

// ' Implements the adjusted z-score
// '
// ' @param y y
// ' @param m m
// ' @param l l
// ' @param s s
// '
// ' @references
// ' http://www.who.int/childgrowth/standards/Chap_7.pdf?ua=1
// '
// ' @noRd
function compute_zscore_adjusted(y: number, m: number, l: number, s: number): number {

  const zscore = compute_zscore(y, m, l, s);
  const SD3pos = calc_sd(3, m, l, s);
  const SD3neg = calc_sd(-3, m, l, s);
  const SD23pos = SD3pos - calc_sd(2, m, l, s);
  const SD23neg = calc_sd(-2, m, l, s) - SD3neg;

  if (zscore > 3) return (3 + ((y - SD3pos) / SD23pos));
  if (zscore < -3) return (-3 + ((y - SD3neg) / SD23neg));
  return zscore;
}

// growthstandards here are the string used to lookup a growth standard
function apply_zscore_and_growthstandards(zscore_fun: zscoreFun, growthstandards: string,
  age_in_days: number, sex: 1 | 2, measure: number): number | undefined {
  if (age_in_days < 0) {
    console.warn(`age_in_days ${age_in_days} is below 0, returning undefined`); 
    return undefined;
  }
  const age_in_days_use = Math.round(age_in_days);

  const standards = lookupGrowthStandards([age_in_days_use, sex], ['age', 'sex'], growthstandards);


  if (standards === undefined) {
    console.warn(`measure ${measure} not found in growth standard ${growthstandards}, returning undefined`)
    return undefined
  }

  const y = measure;
  const m = standards.m;
  const l = standards.l;
  const s = standards.s;
  const zscore: number = zscore_fun(y, m, l, s);
  return round(zscore);
}

function flag_zscore(flag_threshold: number, score_name: string, zscore: number | undefined, valid_zscore: boolean): Record<string, number> | Record<string, undefined> {
  const zname = `z${score_name}`;
  const fname = `f${score_name}`;
  if (!valid_zscore || !zscore) return { [zname]: undefined, [fname]: undefined };
  const fzscore = +(Math.abs(zscore) > flag_threshold);
  return {
    [zname]: zscore,
    [fname]: fzscore,
  };
}

// ' standardise lenhei
// ' if child is <= 730 days, lenhei_unit var should be 'L'.
// ' If lenhei_unit var is 'H', must add 0.7cm to standardise
// ' if child is > 730 days, lenhei_unit var should be 'H'.
// ' If lenhei_unit var is 'L', must subtract 0.7cm to standardise
// ' @noRd
function adjust_lenhei(age_in_days_in: number, measure: string | undefined , lenhei: number): number {
  const age_in_days = Math.round(age_in_days_in);
  if (age_in_days < 731 && measure === 'h') return lenhei + 0.7;
  if (age_in_days >= 731 && measure === 'l') return lenhei - 0.7;
  return lenhei;
}

type anthroZscoreAdjusted = {
  name: string, 
  measure?: number, 
  age_in_days: number, 
  age_in_months: number, 
  sex: 1 | 2, 
  growthstandards: string, 
  flag_threshold: number,
  allowed_age_range?: Array<number>,
  zscore_is_valid?:  boolean,
  zscore_fun?: zscoreFun 
}

function anthro_zscore_adjusted({
  name, measure, age_in_days, age_in_months, sex, growthstandards, flag_threshold,
  allowed_age_range = [0, 1856],
  zscore_is_valid  = true,
  zscore_fun = compute_zscore_adjusted,
}: anthroZscoreAdjusted): ReturnType<typeof flag_zscore> | undefined {
   
  // for all indicators a measure <= 0 should result in zscores being NA
  if (!measure || measure <= 0) return flag_zscore(flag_threshold, name, undefined, false);

  // we convert the input parameter to a data frame and
  // join that with the growthstandards
  // then we have everything to compute the zscores
  const zscore = apply_zscore_and_growthstandards(
    zscore_fun, growthstandards,
    age_in_days, sex, measure,
  );
  if (zscore === undefined) return flag_zscore(flag_threshold, name, undefined, false)

  // we only compute zscores for children age < 60 months
  // the age in months is unrounded
  const valid_age = age_in_months < 60;

  // at last we set certain zscores to NA
  const valid_zscore: boolean = (age_in_days !== undefined
      && age_in_days >= allowed_age_range[0]
      && age_in_days <= allowed_age_range[1]
      && zscore_is_valid
      && valid_age);
  return flag_zscore(flag_threshold, name, zscore, valid_zscore);
}

export {
  compute_zscore,
  compute_zscore_adjusted,
  anthro_zscore_adjusted,
  adjust_lenhei,
  flag_zscore,
  apply_zscore_and_growthstandards
};