package gov.cms.grouper.snf.r2.logic;

import gov.cms.grouper.snf.SnfContext;
import gov.cms.grouper.snf.SnfTables;
import gov.cms.grouper.snf.model.SnfDiagnosisCode;
import gov.cms.grouper.snf.model.enums.Rai300;
import gov.cms.grouper.snf.model.table.BasicRow;
import gov.cms.grouper.snf.model.table.NtaCmgRow;
import gov.cms.grouper.snf.model.table.NtaComorbidityRow;
import gov.cms.grouper.snf.util.ClaimInfo;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=686" class="req">PDPM
 * Payment
 * Component: NTA</a>
 */
public class NtaLogic extends SnfDataVersionImpl<String> {

  public static final Function<SortedSet<String>, Predicate<NtaComorbidityRow>> FeedingFilter =
      (items) -> (row) -> items.equals(row.getMdsItems());
  public static final String ICD10_CODE_HIV = "B20";
  private final Set<String> secondaryDxNtaCategories;
  private final Set<String> assessmentNames;
  private final Supplier<Boolean> hasK0510A2;
  private final Supplier<Integer> k0710A2Value;
  private final Supplier<Integer> k0710B2Value;
  private final Predicate<NtaComorbidityRow> parenteralIvFeeding;
  private final Predicate<NtaComorbidityRow> step1$3AdditionalComorbidities;
  private final ClaimInfo claim;
  public NtaLogic(ClaimInfo claim, List<SnfDiagnosisCode> secondaryDiagnosis) {
    super(claim.getDataVersion());

    this.claim = claim;
    this.assessmentNames =
        claim.getAssessmentNames((item) -> item.isCheck() && !item.getItem().equals("M0300D1"));
    // note on page 688
    this.assessmentNames.addAll(claim
        .getAssessmentNames((item) -> item.getValueInt() > 0 && item.getItem().equals("M0300D1")));

    this.secondaryDxNtaCategories = ClaimInfo.getNtaCategories(secondaryDiagnosis.stream()
        .filter(snfDiagnosisCode -> !snfDiagnosisCode.getValue().equals(ICD10_CODE_HIV))
        .collect(Collectors.toList()));

    this.hasK0510A2 = () -> claim.isCheckedAndNotNull(Rai300.K0510A2);
    this.k0710A2Value = () -> claim.getAssessmentValue(Rai300.K0710A2);
    this.k0710B2Value = () -> claim.getAssessmentValue(Rai300.K0710B2);
    this.parenteralIvFeeding =
        step1$2ParenteralIvFeedingCondition(this.hasK0510A2, this.k0710A2Value, this.k0710B2Value);
    this.step1$3AdditionalComorbidities =
        (row) -> step1$3AdditionalComorbidities(this.assessmentNames, row);
  }

  /**
   * Determine whether the resident meets the criteria for the comorbidity: "Parenteral/IV Feeding –
   * High Intensity" or the comorbidity: "Parenteral/IV Feeding – Low Intensity" <a
   * href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=686" class= "req">step1.2</a>
   */
  public Predicate<NtaComorbidityRow> step1$2ParenteralIvFeedingCondition(
      Supplier<Boolean> hasK0510A2, Supplier<Integer> k0710A2Value,
      Supplier<Integer> k0710B2Value) {
    ParenteralIvFeeding feeding = null;
    if (hasK0510A2.get()) {
      int k0710A2 = k0710A2Value.get();
      if (k0710A2 == 3) {
        feeding = ParenteralIvFeeding.HighIntensity;
      } else if (k0710A2 == 2 && k0710B2Value.get() == 2) {
        feeding = ParenteralIvFeeding.LowIntensity;
      }
    }

    Predicate<NtaComorbidityRow> result = (row) -> Boolean.FALSE;
    if (feeding != null) {
      result = NtaLogic.FeedingFilter.apply(feeding.getItems());
    }

    return SnfContext.trace(feeding, result);
  }

  /**
   * Determine whether the resident has any additional NTA-related comorbidities. <a
   * href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=686" class= "req">step1.3</a>
   */
  public boolean step1$3AdditionalComorbidities(Iterable<String> assessmentNames,
      NtaComorbidityRow row) {
    boolean result = false;
    for (String name : assessmentNames) {
      if (row.getMdsItems().contains(name)
          && !row.getMdsItems().equals(ParenteralIvFeeding.LowIntensity.getItems())
          && !row.getMdsItems().equals(ParenteralIvFeeding.HighIntensity.getItems())) {
        result = true;
        SnfContext.trace(name);
        break;
      }
    }
    return result;
  }

  /**
   * Summarize the resident’s total NTA score from previous steps
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=688" class="req">step
   * 2</a>
   *
   * @return Total NTA score
   */
  public int step2NtaScore(Collection<String> secondaryDxNtaCategories,
      Predicate<NtaComorbidityRow> step1$3AdditionalComorbidities,
      Predicate<NtaComorbidityRow> parenteralIvFeedingFilter) {

    Predicate<NtaComorbidityRow> conditionService = (row) -> secondaryDxNtaCategories.contains(row.getConditionService());

    Predicate<NtaComorbidityRow> conditions = conditionService
        .or(step1$3AdditionalComorbidities)
        .or(parenteralIvFeedingFilter);

    // Secondary Dx must be in secondary nutritional categories
    // Or: assessment has all the items in the assessmentNames
    conditions =
        conditions.and(
            (row) -> ((BiFunction<BasicRow, Integer, Boolean>) BasicRow::isVersion).apply(row, super.getDataVersion()));

    List<Integer> scores = SnfTables.selectAll(SnfTables.ntaComorbidityTableByConditionOfService,
        conditions, (row) -> {
          Integer result = row == null ? (Integer) 0 : row.getPoint();
          Integer point = result;
          return point;
        });
    Integer sum = scores.stream().reduce(0,Integer::sum);

    return SnfContext.trace(sum);

  }

  /**
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=689" class="req">implement
   * step 3</a>
   *
   * @return NTA Case-Mix Group
   */
  public String step3NtaCaseMixGroup(int ntaScore) {
    NtaCmgRow row = SnfTables.ntaCmgTable.values().stream().flatMap(Collection::stream)
        .filter((item) -> item.isVersion(getDataVersion()) && item.isScore(ntaScore)).findFirst()
        .orElse(null);
    String cmg = row.getCmg();

    return SnfContext.trace(cmg);
  }

  /**
   * <a href="doc-files/mds-3.0-rai-manual-v1.17.1_october_2019.pdf#page=686" class="req">Implement
   * PDPM Payment Component: NTA</a>
   */

  @Override
  public String exec() {
    int ntaScore = this.step2NtaScore(this.secondaryDxNtaCategories, step1$3AdditionalComorbidities,
        parenteralIvFeeding);

    String result = this.step3NtaCaseMixGroup(ntaScore);

    return SnfContext.trace(result);

  }

  public Set<String> getSecondaryDxNtaCategories() {
    return secondaryDxNtaCategories;
  }

  public Set<String> getAssessmentNames() {
    return assessmentNames;
  }

  public Supplier<Boolean> getHasK0510A2() {
    return hasK0510A2;
  }

  public Supplier<Integer> getK0710A2Value() {
    return k0710A2Value;
  }

  public Supplier<Integer> getK0710B2Value() {
    return k0710B2Value;
  }

  public Predicate<NtaComorbidityRow> getParenteralIvFeeding() {
    return parenteralIvFeeding;
  }

  public Predicate<NtaComorbidityRow> getStep1$3AdditionalComorbidities() {
    return step1$3AdditionalComorbidities;
  }

  public ClaimInfo getClaim() {
    return claim;
  }

  public enum ParenteralIvFeeding {
    HighIntensity(Arrays.asList(Rai300.K0510A2.name(), Rai300.K0710A2.name())),
    LowIntensity(
        Arrays.asList(Rai300.K0510A2.name(), Rai300.K0710A2.name(), Rai300.K0710B2.name()));

    private final SortedSet<String> items;

    ParenteralIvFeeding(List<String> items) {
      this.items = Collections.unmodifiableSortedSet(new TreeSet<>(items));
    }

    public SortedSet<String> getItems() {
      return items;
    }

  }

}
